Test connection is closed when OAuth token expires

Author: acogoluegnes

.github/workflows/test-pr.yml

@@ -25,7 +25,7 @@ jobs:
       - name: Start broker
         run: ci/start-broker.sh
         env:
-          RABBITMQ_IMAGE: 'pivotalrabbitmq/rabbitmq:amqp-token-renew'
+          RABBITMQ_IMAGE: 'pivotalrabbitmq/rabbitmq:main'
       - name: Start toxiproxy
         run: ci/start-toxiproxy.sh
       - name: Display Java version

ci/start-broker.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 
-RABBITMQ_IMAGE=${RABBITMQ_IMAGE:-pivotalrabbitmq/rabbitmq:amqp-token-renew}
+RABBITMQ_IMAGE=${RABBITMQ_IMAGE:-pivotalrabbitmq/rabbitmq:main}
 
 wait_for_message() {
   while ! docker logs "$1" | grep -q "$2";

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

@@ -331,6 +331,7 @@ TopologyListener createTopologyListener(AmqpConnectionBuilder builder) {
                 "Not recovering connection '{}' for error {}",
                 this.name(),
                 event.failureCause().getMessage());
+            close(ExceptionUtils.convert(ioex));
           }
         };
 
@@ -730,10 +731,10 @@ private void close(Throwable cause) {
         rpcServer.close();
       }
       for (AmqpPublisher publisher : this.publishers) {
-        publisher.close();
+        publisher.close(cause);
       }
       for (AmqpConsumer consumer : this.consumers) {
-        consumer.close();
+        consumer.close(cause);
       }
       try {
         this.dispatchingExecutorService.shutdownNow();

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

@@ -328,7 +328,7 @@ void recoverAfterConnectionFailure() {
     }
   }
 
-  private void close(Throwable cause) {
+  void close(Throwable cause) {
     if (this.closed.compareAndSet(false, true)) {
       this.state(CLOSING, cause);
       this.connection.removeConsumer(this);

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

@@ -205,7 +205,7 @@ private Sender createSender(
     }
   }
 
-  private void close(Throwable cause) {
+  void close(Throwable cause) {
     if (this.closed.compareAndSet(false, true)) {
       this.state(State.CLOSING, cause);
       this.connection.removePublisher(this);

src/test/java/com/rabbitmq/client/amqp/impl/AuthorizationTest.java

@@ -41,7 +41,7 @@ public class AuthorizationTest {
   String name;
 
   @BeforeAll
-  static void initAll() throws Exception {
+  static void initAll() {
     addVhost(VH);
     addUser(USERNAME, PASSWORD);
     setPermissions(USERNAME, VH, "^amqp.*$");

src/test/java/com/rabbitmq/client/amqp/impl/ManagementTest.java

@@ -20,6 +20,7 @@
 import static com.rabbitmq.client.amqp.impl.Assertions.assertThat;
 import static com.rabbitmq.client.amqp.impl.Cli.*;
 import static com.rabbitmq.client.amqp.impl.TestConditions.BrokerVersion.RABBITMQ_4_1_0;
+import static com.rabbitmq.client.amqp.impl.TestUtils.closedOnSecurityExceptionListener;
 import static com.rabbitmq.client.amqp.impl.TestUtils.sync;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -152,13 +153,4 @@ void sessionShouldGetClosedAfterPermissionsChangedAndSetTokenCalled(
       deleteUser(username);
     }
   }
-
-  private static Resource.StateListener closedOnSecurityExceptionListener(Sync sync) {
-    return context -> {
-      if (context.currentState() == Resource.State.CLOSED
-          && context.failureCause() instanceof AmqpException.AmqpSecurityException) {
-        sync.down();
-      }
-    };
-  }
 }

src/test/java/com/rabbitmq/client/amqp/impl/Oauth2Test.java

@@ -17,11 +17,16 @@
 // [email protected].
 package com.rabbitmq.client.amqp.impl;
 
+import static com.rabbitmq.client.amqp.impl.Assertions.assertThat;
+import static com.rabbitmq.client.amqp.impl.TestUtils.*;
+import static java.lang.System.currentTimeMillis;
+import static java.time.Duration.ofMillis;
+import static java.time.Duration.ofSeconds;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
-import com.rabbitmq.client.amqp.AmqpException;
-import com.rabbitmq.client.amqp.Connection;
-import com.rabbitmq.client.amqp.Environment;
+import com.rabbitmq.client.amqp.*;
+import com.rabbitmq.client.amqp.impl.TestUtils.DisabledIfOauth2AuthBackendNotEnabled;
+import com.rabbitmq.client.amqp.impl.TestUtils.Sync;
 import java.time.Duration;
 import java.util.List;
 import org.jose4j.base64url.Base64;
@@ -32,8 +37,10 @@
 import org.jose4j.keys.HmacKey;
 import org.jose4j.lang.JoseException;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
 
 @AmqpTestInfrastructure
+@DisabledIfOauth2AuthBackendNotEnabled
 public class Oauth2Test {
 
   private static final String BASE64_KEY = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH";
@@ -42,19 +49,67 @@ public class Oauth2Test {
 
   @Test
   void expiredTokenShouldFail() throws Exception {
-    String expiredToken = token(System.currentTimeMillis() - 1000);
+    String expiredToken = token(currentTimeMillis() - 1000);
     assertThatThrownBy(
             () -> environment.connectionBuilder().username("").password(expiredToken).build())
         .isInstanceOf(AmqpException.AmqpSecurityException.class);
   }
 
   @Test
   void validTokenShouldSucceed() throws Exception {
-    String validToken = token(System.currentTimeMillis() + Duration.ofMinutes(10).toMillis());
+    String validToken = token(currentTimeMillis() + Duration.ofMinutes(10).toMillis());
     try (Connection ignored =
         environment.connectionBuilder().username("").password(validToken).build()) {}
   }
 
+  @Test
+  void connectionShouldBeClosedWhenTokenExpires(TestInfo info) throws Exception {
+    String q = TestUtils.name(info);
+    long expiry = currentTimeMillis() + ofSeconds(2).toMillis();
+    String token = token(expiry);
+    Sync connectionClosedSync = sync();
+    Connection c =
+        environment
+            .connectionBuilder()
+            .username("")
+            .password(token)
+            .listeners(closedOnSecurityExceptionListener(connectionClosedSync))
+            .build();
+    c.management().queue(q).exclusive(true).declare();
+    Sync publisherClosedSync = sync();
+    Publisher p =
+        c.publisherBuilder()
+            .queue(q)
+            .listeners(closedOnSecurityExceptionListener(publisherClosedSync))
+            .build();
+    Sync consumeSync = sync();
+    Sync consumerClosedSync = sync();
+    c.consumerBuilder()
+        .queue(q)
+        .messageHandler(
+            (ctx, msg) -> {
+              ctx.accept();
+              consumeSync.down();
+            })
+        .listeners(closedOnSecurityExceptionListener(consumerClosedSync))
+        .build();
+    p.publish(p.message(), ctx -> {});
+    assertThat(consumeSync).completes();
+    long newExpiry = currentTimeMillis() + ofSeconds(3).toMillis();
+    token = token(newExpiry);
+    ((AmqpManagement) c.management()).setToken(token);
+    // wait for the first token to expire
+    waitAtMost(() -> currentTimeMillis() > expiry + ofMillis(500).toMillis());
+    // the connection should still work thanks to the new token
+    consumeSync.reset();
+    p.publish(p.message(), ctx -> {});
+    assertThat(consumeSync).completes();
+    waitAtMost(() -> currentTimeMillis() > newExpiry);
+    assertThat(connectionClosedSync).completes();
+    assertThat(publisherClosedSync).completes();
+    assertThat(consumerClosedSync).completes();
+  }
+
   private static String token(long expirationTime) throws JoseException {
     JwtClaims claims = new JwtClaims();
     claims.setIssuer("unit_test");

src/test/java/com/rabbitmq/client/amqp/impl/TestUtils.java

@@ -21,6 +21,8 @@
 import static java.util.Collections.singletonMap;
 import static org.assertj.core.api.Assertions.fail;
 
+import com.rabbitmq.client.amqp.AmqpException;
+import com.rabbitmq.client.amqp.Resource;
 import eu.rekawek.toxiproxy.Proxy;
 import eu.rekawek.toxiproxy.ToxiproxyClient;
 import java.io.IOException;
@@ -83,6 +85,15 @@ static <T> T waitUntilStable(Supplier<T> call, Duration waitTime) {
     return null;
   }
 
+  static Resource.StateListener closedOnSecurityExceptionListener(Sync sync) {
+    return context -> {
+      if (context.currentState() == Resource.State.CLOSED
+          && context.failureCause() instanceof AmqpException.AmqpSecurityException) {
+        sync.down();
+      }
+    };
+  }
+
   @FunctionalInterface
   public interface CallableBooleanSupplier {
     boolean getAsBoolean() throws Exception;
@@ -404,6 +415,16 @@ private static class DisabledIfAuthMechanismSslNotEnabledCondition
     }
   }
 
+  private static class DisabledIfOauth2AuthBackendNotEnabledCondition
+      extends DisabledIfPluginNotEnabledCondition {
+
+    DisabledIfOauth2AuthBackendNotEnabledCondition() {
+      super(
+          "OAuth2 authentication backend",
+          output -> output.contains("rabbitmq_auth_backend_oauth2"));
+    }
+  }
+
   static class DisabledIfNotClusterCondition implements ExecutionCondition {
 
     private static final String KEY = "isCluster";
@@ -463,6 +484,12 @@ public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext con
   @ExtendWith(DisabledIfAuthMechanismSslNotEnabledCondition.class)
   @interface DisabledIfAuthMechanismSslNotEnabled {}
 
+  @Target({ElementType.TYPE, ElementType.METHOD})
+  @Retention(RetentionPolicy.RUNTIME)
+  @Documented
+  @ExtendWith(DisabledIfOauth2AuthBackendNotEnabledCondition.class)
+  @interface DisabledIfOauth2AuthBackendNotEnabled {}
+
   @Target({ElementType.TYPE, ElementType.METHOD})
   @Retention(RetentionPolicy.RUNTIME)
   @Documented