Start support for token refresh

Author: acogoluegnes

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

@@ -76,7 +76,7 @@ final class AmqpConnection extends ResourceBase implements Connection {
   private final Lock instanceLock = new ReentrantLock();
   private final boolean filterExpressionsSupported, setTokenSupported;
   private volatile ExecutorService dispatchingExecutorService;
-  private final Credentials credentials;
+  private final Credentials.Registration credentialsRegistration;
 
   AmqpConnection(AmqpConnectionBuilder builder) {
     super(builder.listeners());
@@ -119,9 +119,14 @@ final class AmqpConnection extends ResourceBase implements Connection {
         instanceof UsernamePasswordCredentialsProvider) {
       UsernamePasswordCredentialsProvider credentialsProvider =
           (UsernamePasswordCredentialsProvider) connectionSettings.credentialsProvider();
-      this.credentials = new UsernamePasswordCredentials(credentialsProvider);
+      Credentials credentials = new UsernamePasswordCredentials(credentialsProvider);
+      this.credentialsRegistration =
+          credentials.register(
+              (u, p) -> {
+                ((AmqpManagement) management()).setToken(p);
+              });
     } else {
-      this.credentials = callback -> {};
+      this.credentialsRegistration = Credentials.NO_OP.register((u, p) -> {});
     }
     LOGGER.debug("Opening native connection for connection '{}'...", this.name());
     NativeConnectionWrapper ncw =
@@ -199,7 +204,7 @@ private NativeConnectionWrapper connect(
       List<Address> addresses) {
 
     ConnectionOptions connectionOptions = new ConnectionOptions();
-    credentials.configure(new TokenConnectionCallback(connectionOptions));
+    credentialsRegistration.connect(new TokenConnectionCallback(connectionOptions));
     connectionOptions.virtualHost("vhost:" + connectionSettings.virtualHost());
     connectionOptions.saslOptions().addAllowedMechanism(connectionSettings.saslMechanism());
     connectionOptions.idleTimeout(
@@ -726,6 +731,7 @@ long id() {
   private void close(Throwable cause) {
     if (this.closed.compareAndSet(false, true)) {
       this.state(CLOSING, cause);
+      this.credentialsRegistration.unregister();
       this.environment.removeConnection(this);
       if (this.topologyListener instanceof AutoCloseable) {
         try {

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

@@ -19,12 +19,43 @@
 
 interface Credentials {
 
-  void configure(ConnectionCallback callback);
+  Credentials NO_OP = new NoOpCredentials();
+
+  Registration register(RefreshCallback refreshCallback);
+
+  interface Registration {
+
+    void connect(ConnectionCallback callback);
+
+    void unregister();
+  }
 
   interface ConnectionCallback {
 
     ConnectionCallback username(String username);
 
     ConnectionCallback password(String password);
   }
+
+  interface RefreshCallback {
+
+    void refresh(String username, String password);
+  }
+
+  class NoOpCredentials implements Credentials {
+
+    @Override
+    public Registration register(RefreshCallback refreshCallback) {
+      return new NoOpRegistration();
+    }
+  }
+
+  class NoOpRegistration implements Registration {
+
+    @Override
+    public void connect(ConnectionCallback callback) {}
+
+    @Override
+    public void unregister() {}
+  }
 }

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

@@ -19,21 +19,18 @@
 
 import com.rabbitmq.client.amqp.UsernamePasswordCredentialsProvider;
 import com.rabbitmq.client.amqp.oauth.Token;
-import com.rabbitmq.client.amqp.oauth.TokenParser;
 import com.rabbitmq.client.amqp.oauth.TokenRequester;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 
 final class OAuthCredentialsProvider implements UsernamePasswordCredentialsProvider {
 
   private final TokenRequester requester;
-  private final TokenParser parser;
   private volatile Token token;
   private final Lock lock = new ReentrantLock();
 
-  OAuthCredentialsProvider(TokenRequester requester, TokenParser parser) {
+  OAuthCredentialsProvider(TokenRequester requester) {
     this.requester = requester;
-    this.parser = parser;
   }
 
   @Override
@@ -46,7 +43,7 @@ public String getPassword() {
     lock.lock();
     try {
       if (token == null || token.expirationTime() < System.currentTimeMillis()) {
-        this.token = this.parser.parse(this.requester.request());
+        this.token = this.requester.request();
       }
     } finally {
       lock.unlock();

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

@@ -0,0 +1,184 @@
+// 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.oauth.Token;
+import com.rabbitmq.client.amqp.oauth.TokenRequester;
+import java.time.Duration;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+final class TokenCredentials implements Credentials {
+
+  private final TokenRequester requester;
+  private final ScheduledExecutorService scheduledExecutorService;
+  private volatile Token token;
+  private final Lock lock = new ReentrantLock();
+  private final Map<Long, RegistrationImpl> registrations = new ConcurrentHashMap<>();
+  private final AtomicLong registrationSequence = new AtomicLong(0);
+  private final AtomicBoolean schedulingRenewal = new AtomicBoolean(false);
+  private volatile ScheduledFuture<?> renewalTask;
+
+  TokenCredentials(TokenRequester requester, ScheduledExecutorService scheduledExecutorService) {
+    this.requester = requester;
+    this.scheduledExecutorService = scheduledExecutorService;
+  }
+
+  private void lock() {
+    this.lock.lock();
+  }
+
+  private void unlock() {
+    this.lock.unlock();
+  }
+
+  private boolean expiresSoon(Token t) {
+    return t.expirationTime() < System.currentTimeMillis() - 20_000;
+  }
+
+  private Duration delayBeforeTokenRenewal(Token token) {
+    long expiresIn = token.expirationTime() - System.currentTimeMillis();
+    long delay = (long) (expiresIn * 0.8);
+    return Duration.ofMillis(delay);
+  }
+
+  private Token getToken() {
+    return requester.request();
+  }
+
+  @Override
+  public Registration register(RefreshCallback refreshCallback) {
+    Long id = this.registrationSequence.getAndIncrement();
+    RegistrationImpl registration = new RegistrationImpl(id, refreshCallback);
+    this.registrations.put(id, registration);
+    return registration;
+  }
+
+  private void refresh() {
+    this.scheduledExecutorService.execute(
+        () -> {
+          for (RegistrationImpl registration : this.registrations.values()) {
+            if (!registration.isClosed() && !this.token.equals(registration.registrationToken)) {
+              // the registration does not have the new token yet
+              registration.refreshCallback.refresh("", this.token.value());
+              registration.registrationToken = this.token;
+            }
+          }
+        });
+  }
+
+  private void token(Token t) {
+    if (!t.equals(this.token)) {
+      this.token = t;
+      if (this.schedulingRenewal.compareAndSet(false, true)) {
+        if (this.renewalTask != null) {
+          this.renewalTask.cancel(false);
+        }
+        Duration delay = delayBeforeTokenRenewal(t);
+        if (delay.isZero() || delay.isNegative()) {
+          delay = Duration.ofSeconds(1);
+        }
+        // TODO check delay is > 0, schedule 1 second later at least
+        this.renewalTask =
+            this.scheduledExecutorService.schedule(
+                () -> {
+                  Token previousToken = this.token;
+                  this.lock();
+                  try {
+                    if (this.token.equals(previousToken)) {
+                      Token newToken = getToken();
+                      token(newToken);
+                      refresh();
+                    }
+                  } finally {
+                    unlock();
+                  }
+                },
+                delay.toMillis(),
+                TimeUnit.MILLISECONDS);
+        this.schedulingRenewal.set(false);
+      }
+    }
+  }
+
+  private final class RegistrationImpl implements Registration {
+
+    private final Long id;
+    private final RefreshCallback refreshCallback;
+    private volatile Token registrationToken;
+    private final AtomicBoolean closed = new AtomicBoolean(false);
+
+    private RegistrationImpl(Long id, RefreshCallback refreshCallback) {
+      this.id = id;
+      this.refreshCallback = refreshCallback;
+    }
+
+    @Override
+    public void connect(ConnectionCallback callback) {
+      boolean shouldRefresh = false;
+      lock();
+      try {
+        if (token == null) {
+          token(getToken());
+        } else if (expiresSoon(token)) {
+          shouldRefresh = true;
+          token(getToken());
+        }
+        this.registrationToken = token;
+      } finally {
+        unlock();
+      }
+
+      callback.username("").password(this.registrationToken.value());
+      if (shouldRefresh) {
+        refresh();
+      }
+    }
+
+    @Override
+    public void unregister() {
+      if (this.closed.compareAndSet(false, true)) {
+        registrations.remove(this.id);
+      }
+    }
+
+    private boolean isClosed() {
+      return this.closed.get();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (o == null || getClass() != o.getClass()) return false;
+      RegistrationImpl that = (RegistrationImpl) o;
+      return Objects.equals(id, that.id);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(id);
+    }
+  }
+}

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

@@ -28,7 +28,18 @@ final class UsernamePasswordCredentials implements Credentials {
   }
 
   @Override
-  public void configure(ConnectionCallback callback) {
-    callback.username(this.provider.getUsername()).password(this.provider.getPassword());
+  public Registration register(RefreshCallback refreshCallback) {
+    return new RegistrationImpl();
+  }
+
+  private final class RegistrationImpl implements Registration {
+
+    @Override
+    public void connect(ConnectionCallback callback) {
+      callback.username(provider.getUsername()).password(provider.getPassword());
+    }
+
+    @Override
+    public void unregister() {}
   }
 }

src/main/java/com/rabbitmq/client/amqp/oauth/HttpTokenRequester.java

@@ -51,6 +51,8 @@ public final class HttpTokenRequester implements TokenRequester {
   private final HttpClient client;
   private final Consumer<HttpRequest.Builder> requestBuilderConsumer;
 
+  private TokenParser parser;
+
   public HttpTokenRequester(
       String tokenEndpointUri,
       String clientId,
@@ -60,7 +62,8 @@ public HttpTokenRequester(
       HostnameVerifier hostnameVerifier,
       SSLSocketFactory sslSocketFactory,
       Consumer<HttpClient.Builder> clientBuilderConsumer,
-      Consumer<HttpRequest.Builder> requestBuilderConsumer) {
+      Consumer<HttpRequest.Builder> requestBuilderConsumer,
+      TokenParser parser) {
     try {
       this.tokenEndpointUri = new URI(tokenEndpointUri);
     } catch (URISyntaxException e) {
@@ -72,6 +75,7 @@ public HttpTokenRequester(
     this.parameters = Map.copyOf(parameters);
     this.hostnameVerifier = hostnameVerifier;
     this.sslSocketFactory = sslSocketFactory;
+    this.parser = parser;
     if (requestBuilderConsumer == null) {
       this.requestBuilderConsumer =
           requestBuilder ->
@@ -95,7 +99,7 @@ public HttpTokenRequester(
   }
 
   @Override
-  public String request() {
+  public Token request() {
     StringBuilder urlParameters = new StringBuilder();
     encode(urlParameters, "grant_type", grantType);
     for (Map.Entry<String, String> parameter : parameters.entrySet()) {
@@ -117,7 +121,7 @@ public String request() {
           this.client.send(request, HttpResponse.BodyHandlers.ofString(UTF_8));
       checkStatusCode(response.statusCode());
       checkContentType(response.headers().firstValue("content-type").orElse(null));
-      return response.body();
+      return this.parser.parse(response.body());
     } catch (IOException e) {
       throw new OAuthException("Error while retrieving OAuth 2 token", e);
     } catch (InterruptedException e) {

src/main/java/com/rabbitmq/client/amqp/oauth/TokenRequester.java

@@ -19,5 +19,5 @@
 
 public interface TokenRequester {
 
-  String request();
+  Token request();
 }