Close producer and consumer on security exception

The session can be closed after permissions are re-evaluated, e.g. when
a token is renewed.

Author: acogoluegnes

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

@@ -19,6 +19,7 @@
 
 import static com.rabbitmq.client.amqp.Resource.State.*;
 import static com.rabbitmq.client.amqp.impl.AmqpConsumerBuilder.*;
+import static com.rabbitmq.client.amqp.impl.ExceptionUtils.*;
 import static java.time.Duration.ofSeconds;
 import static java.util.Optional.ofNullable;
 
@@ -71,7 +72,7 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
   private final Runnable replenishCreditOperation = this::replenishCreditIfNeeded;
   private final ExecutorService dispatchingExecutorService;
   private final java.util.function.Consumer<Delivery> nativeHandler;
-  private final java.util.function.Consumer<ClientException> nativeReceiverCloseHandler;
+  private final java.util.function.Consumer<ClientException> nativeCloseHandler;
   // native receiver internal state, accessed only in the native executor/scheduler
   private ProtonReceiver protonReceiver;
   private volatile Scheduler protonExecutor;
@@ -101,7 +102,7 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
 
     this.dispatchingExecutorService = connection.dispatchingExecutorService();
     this.nativeHandler = createNativeHandler(messageHandler);
-    this.nativeReceiverCloseHandler =
+    this.nativeCloseHandler =
         e ->
             this.dispatchingExecutorService.submit(
                 () -> {
@@ -116,7 +117,7 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
             this.filters,
             this.subscriptionListener,
             this.nativeHandler,
-            this.nativeReceiverCloseHandler);
+            this.nativeCloseHandler);
     this.initStateFromNativeReceiver(this.nativeReceiver);
     this.metricsCollector = this.connection.metricsCollector();
     try {
@@ -277,8 +278,8 @@ private Runnable createReceiveTask(Receiver receiver, MessageHandler messageHand
             messageHandler.handle(context, message);
           }
         }
-      } catch (ClientLinkRemotelyClosedException e) {
-        if (ExceptionUtils.notFound(e) || ExceptionUtils.resourceDeleted(e)) {
+      } catch (ClientLinkRemotelyClosedException | ClientSessionRemotelyClosedException e) {
+        if (notFound(e) || resourceDeleted(e) || unauthorizedAccess(e)) {
           this.close(ExceptionUtils.convert(e));
         }
       } catch (ClientConnectionRemotelyClosedException e) {
@@ -304,7 +305,7 @@ void recoverAfterConnectionFailure() {
                     this.filters,
                     this.subscriptionListener,
                     this.nativeHandler,
-                    this.nativeReceiverCloseHandler),
+                    this.nativeCloseHandler),
             e -> {
               boolean shouldRetry =
                   e instanceof AmqpException.AmqpResourceClosedException
@@ -533,13 +534,6 @@ private void handleException(Exception ex, String operation) {
   }
 
   private static boolean maybeCloseConsumerOnException(AmqpConsumer consumer, Exception ex) {
-    if (ex instanceof ClientLinkRemotelyClosedException) {
-      ClientLinkRemotelyClosedException e = (ClientLinkRemotelyClosedException) ex;
-      if (ExceptionUtils.notFound(e) || ExceptionUtils.resourceDeleted(e)) {
-        consumer.close(ExceptionUtils.convert(e));
-        return true;
-      }
-    }
-    return false;
+    return ExceptionUtils.maybeCloseConsumerOnException(consumer::close, ex);
   }
 }

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

@@ -27,6 +27,7 @@
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import org.apache.qpid.protonj2.client.*;
 import org.apache.qpid.protonj2.client.exceptions.ClientException;
@@ -53,6 +54,8 @@ final class AmqpPublisher extends ResourceBase implements Publisher {
   private final Duration publishTimeout;
   private final SessionHandler sessionHandler;
   private volatile ObservationCollector.ConnectionInfo connectionInfo;
+  private final ExecutorService dispatchingExecutorService;
+  private final java.util.function.Consumer<ClientException> nativeCloseHandler;
 
   AmqpPublisher(AmqpPublisherBuilder builder) {
     super(builder.listeners());
@@ -63,7 +66,17 @@ final class AmqpPublisher extends ResourceBase implements Publisher {
     this.connection = builder.connection();
     this.publishTimeout = builder.publishTimeout();
     this.sessionHandler = this.connection.createSessionHandler();
-    this.sender = this.createSender(sessionHandler.session(), this.address, this.publishTimeout);
+    this.dispatchingExecutorService = connection.dispatchingExecutorService();
+    this.nativeCloseHandler =
+        e ->
+            this.dispatchingExecutorService.submit(
+                () -> {
+                  // get result to make spotbugs happy
+                  boolean ignored = maybeCloseConsumerOnException(this, e);
+                });
+    this.sender =
+        this.createSender(
+            sessionHandler.session(), this.address, this.publishTimeout, this.nativeCloseHandler);
     this.metricsCollector = this.connection.metricsCollector();
     this.observationCollector = this.connection.observationCollector();
     this.state(OPEN);
@@ -154,7 +167,11 @@ private static MetricsCollector.PublishDisposition mapToPublishDisposition(Statu
   void recoverAfterConnectionFailure() {
     this.connectionInfo = new Utils.ObservationConnectionInfo(this.connection.connectionAddress());
     this.sender =
-        this.createSender(this.sessionHandler.sessionNoCheck(), this.address, this.publishTimeout);
+        this.createSender(
+            this.sessionHandler.sessionNoCheck(),
+            this.address,
+            this.publishTimeout,
+            this.nativeCloseHandler);
   }
 
   @Override
@@ -164,14 +181,19 @@ public void close() {
 
   // internal API
 
-  private Sender createSender(Session session, String address, Duration publishTimeout) {
+  private Sender createSender(
+      Session session,
+      String address,
+      Duration publishTimeout,
+      Consumer<ClientException> nativeCloseHandler) {
     SenderOptions senderOptions =
         new SenderOptions()
             .deliveryMode(DeliveryMode.AT_LEAST_ONCE)
             .sendTimeout(
                 publishTimeout.isNegative()
                     ? ConnectionOptions.INFINITE
-                    : publishTimeout.toMillis());
+                    : publishTimeout.toMillis())
+            .closeHandler(nativeCloseHandler);
     try {
       Sender s =
           address == null
@@ -198,6 +220,10 @@ private void close(Throwable cause) {
     }
   }
 
+  private static boolean maybeCloseConsumerOnException(AmqpPublisher publisher, Exception ex) {
+    return ExceptionUtils.maybeCloseConsumerOnException(publisher::close, ex);
+  }
+
   private static class DefaultContext implements Publisher.Context {
 
     private final Message message;

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

@@ -20,6 +20,7 @@
 import com.rabbitmq.client.amqp.AmqpException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.function.Consumer;
 import javax.net.ssl.SSLException;
 import org.apache.qpid.protonj2.client.ErrorCondition;
 import org.apache.qpid.protonj2.client.exceptions.*;
@@ -67,20 +68,14 @@ static AmqpException convert(ClientException e, String format, Object... args) {
       return new AmqpException.AmqpSecurityException(message, e);
     } else if (isNetworkError(e)) {
       return new AmqpException.AmqpConnectionException(e.getMessage(), e);
-    } else if (e instanceof ClientSessionRemotelyClosedException) {
+    } else if (e instanceof ClientSessionRemotelyClosedException
+        || e instanceof ClientLinkRemotelyClosedException) {
       ErrorCondition errorCondition =
-          ((ClientSessionRemotelyClosedException) e).getErrorCondition();
+          ((ClientResourceRemotelyClosedException) e).getErrorCondition();
       if (isUnauthorizedAccess(errorCondition)) {
         return new AmqpException.AmqpSecurityException(e.getMessage(), e);
       } else if (isNotFound(errorCondition)) {
         return new AmqpException.AmqpEntityDoesNotExistException(e.getMessage(), e);
-      } else {
-        return new AmqpException.AmqpResourceClosedException(e.getMessage(), e);
-      }
-    } else if (e instanceof ClientLinkRemotelyClosedException) {
-      ErrorCondition errorCondition = ((ClientLinkRemotelyClosedException) e).getErrorCondition();
-      if (isNotFound(errorCondition)) {
-        return new AmqpException.AmqpEntityDoesNotExistException(e.getMessage(), e);
       } else if (isResourceDeleted(errorCondition)) {
         return new AmqpException.AmqpEntityDoesNotExistException(e.getMessage(), e);
       } else {
@@ -109,6 +104,10 @@ static boolean notFound(ClientResourceRemotelyClosedException e) {
         && "amqp:not-found".equals(e.getErrorCondition().condition());
   }
 
+  static boolean unauthorizedAccess(ClientResourceRemotelyClosedException e) {
+    return isUnauthorizedAccess(e.getErrorCondition());
+  }
+
   private static boolean isUnauthorizedAccess(ErrorCondition errorCondition) {
     return errorConditionEquals(errorCondition, ERROR_UNAUTHORIZED_ACCESS);
   }
@@ -135,4 +134,16 @@ private static boolean isNetworkError(ClientException e) {
     }
     return false;
   }
+
+  static boolean maybeCloseConsumerOnException(Consumer<Throwable> closing, Exception ex) {
+    if (ex instanceof ClientLinkRemotelyClosedException
+        || ex instanceof ClientSessionRemotelyClosedException) {
+      ClientResourceRemotelyClosedException e = (ClientResourceRemotelyClosedException) ex;
+      if (notFound(e) || resourceDeleted(e) || unauthorizedAccess(e)) {
+        closing.accept(ExceptionUtils.convert(e));
+        return true;
+      }
+    }
+    return false;
+  }
 }

src/main/qpid/org/apache/qpid/protonj2/client/SenderOptions.java

@@ -17,8 +17,10 @@
 package org.apache.qpid.protonj2.client;
 
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
+import org.apache.qpid.protonj2.client.exceptions.ClientException;
 import org.apache.qpid.protonj2.client.exceptions.ClientSendTimedOutException;
 import org.apache.qpid.protonj2.engine.DeliveryTagGenerator;
 
@@ -28,6 +30,7 @@
 public class SenderOptions extends LinkOptions<SenderOptions> implements Cloneable {
 
     private long sendTimeout = ConnectionOptions.DEFAULT_SEND_TIMEOUT;
+    private Consumer<ClientException> closeHandler;
 
     private Supplier<DeliveryTagGenerator> tagGeneratorSupplier;
 
@@ -89,6 +92,26 @@ public SenderOptions sendTimeout(long timeout, TimeUnit units) {
         return this;
     }
 
+    /**
+     * Callback when the sender is closed / shut down.
+     *
+     * @param closeHandler close / shutdown handler
+     * @return this {@link SenderOptions} instance.
+     */
+    public SenderOptions closeHandler(Consumer<ClientException> closeHandler) {
+        this.closeHandler = closeHandler;
+        return this;
+    }
+
+    /**
+     * Configured close / shutdown handler.
+     *
+     * @return the configured handler
+     */
+    public Consumer<ClientException> closeHandler() {
+        return this.closeHandler;
+    }
+
     @Override
     public SenderOptions clone() {
         return copyInto(new SenderOptions());

src/main/qpid/org/apache/qpid/protonj2/client/impl/ClientSender.java

@@ -21,6 +21,7 @@
 import java.util.Map;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 import org.apache.qpid.protonj2.buffer.ProtonBuffer;
 import org.apache.qpid.protonj2.buffer.ProtonBufferAllocator;
@@ -49,11 +50,13 @@ public final class ClientSender extends ClientSenderLinkType<Sender> implements
 
     private final Deque<ClientOutgoingEnvelope> blocked = new ArrayDeque<>();
     private final SenderOptions options;
+    private final Consumer<ClientException> closeHandler;
 
     ClientSender(ClientSession session, SenderOptions options, String senderId, org.apache.qpid.protonj2.engine.Sender protonSender) {
         super(session, senderId, options, protonSender);
 
         this.options = new SenderOptions(options);
+        this.closeHandler = options.closeHandler() == null ? e ->  { } : options.closeHandler();
     }
 
     @Override
@@ -208,6 +211,11 @@ protected void linkSpecificCleanupHandler(ClientException failureCause) {
         } else {
             failPendingUnsettledAndBlockedSends(new ClientResourceRemotelyClosedException("The sender link has closed"));
         }
+        try {
+            this.closeHandler.accept(failureCause);
+        } catch (Exception e) {
+            LOG.warn("Error in close handler: {}", e.getMessage());
+        }
     }
 
     private void failPendingUnsettledAndBlockedSends(ClientException cause) {