JAMES-4210 Improve SaslExchange cleanup

protocols/api/src/main/java/org/apache/james/protocols/api/sasl/SaslExchange.java

@@ -34,12 +34,9 @@ public interface SaslExchange extends AutoCloseable {
     SaslStep onResponse(byte[] clientResponse);
 
     /**
-     * Aborts the exchange after a client cancellation or protocol-level failure, and releases associated resources.
+     * Releases resources associated with the exchange after any terminal outcome, including success,
+     * failure, cancellation, disconnect, timeout, or protocol error.
      */
-    default void abort() {
-        close();
-    }
-
     @Override
     void close();
 }

protocols/api/src/test/java/org/apache/james/protocols/api/sasl/SaslMechanismContractTest.java

@@ -70,7 +70,6 @@ public SaslExchange start(SaslInitialRequest request, SaslAuthenticator authenti
 
     private static class FixedStepExchange implements SaslExchange {
         private final SaslStep firstStep;
-        private boolean aborted;
         private boolean closed;
 
         private FixedStepExchange(SaslStep firstStep) {
@@ -87,11 +86,6 @@ public SaslStep onResponse(byte[] clientResponse) {
             return firstStep;
         }
 
-        @Override
-        public void abort() {
-            aborted = true;
-        }
-
         @Override
         public void close() {
             closed = true;
@@ -133,10 +127,6 @@ public SaslStep onResponse(byte[] clientResponse) {
             return new SaslStep.Failure(SaslFailure.invalidCredentials(AUTHENTICATION_ID, Optional.empty(), "rejected"));
         }
 
-        @Override
-        public void abort() {
-        }
-
         @Override
         public void close() {
         }
@@ -226,16 +216,14 @@ void saslStepsShouldDefensivelyCopyPayloads() {
     }
 
     @Test
-    void exchangeShouldExposeAbortAndCloseLifecycle() {
+    void exchangeShouldExposeCloseLifecycle() {
         // GIVEN an active exchange
         FixedStepExchange exchange = new FixedStepExchange(new SaslStep.Failure(SaslFailure.malformed("failure")));
 
-        // WHEN the protocol aborts and then closes it
-        exchange.abort();
+        // WHEN the protocol terminates it
         exchange.close();
 
-        // THEN mechanisms can observe both lifecycle events
-        assertThat(exchange.aborted).isTrue();
+        // THEN the mechanism can release its resources
         assertThat(exchange.closed).isTrue();
     }
 

protocols/imap/src/main/java/org/apache/james/imap/api/process/ImapSaslExchangeTracker.java

@@ -0,0 +1,109 @@
+/****************************************************************
+ * Licensed to the Apache Software Foundation (ASF) under one   *
+ * or more contributor license agreements.  See the NOTICE file *
+ * distributed with this work for additional information        *
+ * regarding copyright ownership.  The ASF licenses this file   *
+ * to you 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.                                           *
+ ****************************************************************/
+
+package org.apache.james.imap.api.process;
+
+import org.apache.james.protocols.api.sasl.SaslExchange;
+
+/**
+ * Owns the active SASL exchange for an IMAP session. Once closed on disconnect,
+ * the tracker stays attached and rejects delayed asynchronous registrations.
+ */
+public class ImapSaslExchangeTracker {
+    private static final String ATTRIBUTE_KEY = ImapSaslExchangeTracker.class.getName();
+
+    public static ImapSaslExchangeTracker forSession(ImapSession session) {
+        // Make tracker initialization atomic with disconnect sealing for this session.
+        synchronized (session) {
+            Object value = session.getAttribute(ATTRIBUTE_KEY);
+            if (value instanceof ImapSaslExchangeTracker tracker) {
+                return tracker;
+            }
+
+            ImapSaslExchangeTracker tracker = new ImapSaslExchangeTracker();
+            session.setAttribute(ATTRIBUTE_KEY, tracker);
+            return tracker;
+        }
+    }
+
+    public static void closeForSession(ImapSession session) {
+        forSession(session).close();
+    }
+
+    private static IllegalStateException closeRejectedExchange(SaslExchange exchange) {
+        IllegalStateException failure = new IllegalStateException("IMAP SASL exchange cannot be registered");
+        try {
+            exchange.close();
+        } catch (RuntimeException e) {
+            failure.addSuppressed(e);
+        }
+        return failure;
+    }
+
+    private SaslExchange activeExchange;
+    private boolean closed;
+
+    private ImapSaslExchangeTracker() {
+    }
+
+    public SaslExchange register(SaslExchange exchange) {
+        if (tryRegister(exchange)) {
+            return exchange;
+        }
+        throw closeRejectedExchange(exchange);
+    }
+
+    private synchronized boolean tryRegister(SaslExchange exchange) {
+        if (closed || activeExchange != null) {
+            return false;
+        }
+        activeExchange = exchange;
+        return true;
+    }
+
+    public void closeExchange(SaslExchange exchange) {
+        if (release(exchange)) {
+            exchange.close();
+        }
+    }
+
+    public void close() {
+        SaslExchange exchange;
+        synchronized (this) {
+            if (closed) {
+                return;
+            }
+            closed = true;
+            exchange = activeExchange;
+            activeExchange = null;
+        }
+
+        if (exchange != null) {
+            exchange.close();
+        }
+    }
+
+    private synchronized boolean release(SaslExchange exchange) {
+        if (activeExchange != exchange) {
+            return false;
+        }
+        activeExchange = null;
+        return true;
+    }
+}

protocols/imap/src/main/java/org/apache/james/imap/processor/AuthenticateProcessor.java

@@ -29,6 +29,7 @@
 import org.apache.james.imap.api.display.HumanReadableText;
 import org.apache.james.imap.api.message.Capability;
 import org.apache.james.imap.api.message.response.StatusResponseFactory;
+import org.apache.james.imap.api.process.ImapSaslExchangeTracker;
 import org.apache.james.imap.api.process.ImapSession;
 import org.apache.james.imap.main.PathConverter;
 import org.apache.james.imap.message.request.AuthenticateRequest;
@@ -97,8 +98,9 @@ protected void processRequest(AuthenticateRequest request, ImapSession session,
         try {
             SaslInitialRequest initialRequest = SaslCodec.initialRequest(request.getAuthType(), initialClientResponse(request));
             SaslAuthenticator authenticator = jamesSaslAuthenticator.withExtraAuthorizator(withAdminUsers());
-            SaslExchange exchange = mechanism.get().start(initialRequest, authenticator);
-            handleFirstStep(exchange, firstStep(exchange), session, request, responder);
+            SaslExchange exchange = ImapSaslExchangeTracker.forSession(session)
+                .register(mechanism.get().start(initialRequest, authenticator));
+            handleFirstStep(exchange, firstStep(exchange, session), session, request, responder);
         } catch (IllegalArgumentException e) {
             LOGGER.info("Invalid syntax in AUTHENTICATE initial client response", e);
             authFailure(session, request, responder, HumanReadableText.AUTHENTICATION_FAILED, Optional.empty(),
@@ -136,11 +138,11 @@ private Optional<String> initialClientResponse(AuthenticateRequest request) {
         return Optional.empty();
     }
 
-    private SaslStep firstStep(SaslExchange exchange) {
+    private SaslStep firstStep(SaslExchange exchange, ImapSession session) {
         try {
             return exchange.firstStep();
         } catch (RuntimeException e) {
-            exchange.close();
+            ImapSaslExchangeTracker.forSession(session).closeExchange(exchange);
             throw e;
         }
     }
@@ -195,14 +197,14 @@ private void pushContinuationHandler(SaslExchange exchange, ImapSession session,
                 .subscribeOn(ReactorUtils.BLOCKING_CALL_WRAPPER)
                 .then());
         } catch (RuntimeException e) {
-            exchange.close();
+            ImapSaslExchangeTracker.forSession(session).closeExchange(exchange);
             throw e;
         }
     }
 
     private void handleContinuationLine(SaslExchange exchange, ImapSession session, AuthenticateRequest request, Responder responder, byte[] data) {
         if (isAbort(exchange, session, data)) {
-            abortActiveContinuation(exchange, session);
+            closeActiveContinuation(exchange, session);
             no(request, responder, HumanReadableText.AUTHENTICATION_FAILED);
             responder.flush();
             return;
@@ -281,23 +283,15 @@ private void closeActiveContinuation(SaslExchange exchange, ImapSession session)
         try {
             session.popLineHandler();
         } finally {
-            exchange.close();
-        }
-    }
-
-    private void abortActiveContinuation(SaslExchange exchange, ImapSession session) {
-        try {
-            session.popLineHandler();
-        } finally {
-            exchange.abort();
+            ImapSaslExchangeTracker.forSession(session).closeExchange(exchange);
         }
     }
 
     private void popActiveContinuation(SaslExchange exchange, ImapSession session) {
         try {
             session.popLineHandler();
         } catch (RuntimeException e) {
-            exchange.close();
+            ImapSaslExchangeTracker.forSession(session).closeExchange(exchange);
             throw e;
         }
     }
@@ -318,15 +312,15 @@ private void pushSuccessDataAcknowledgementHandler(SaslExchange exchange, SaslSt
                 .subscribeOn(ReactorUtils.BLOCKING_CALL_WRAPPER)
                 .then());
         } catch (RuntimeException e) {
-            exchange.close();
+            ImapSaslExchangeTracker.forSession(session).closeExchange(exchange);
             throw e;
         }
     }
 
     private void handleSuccessDataAcknowledgement(SaslExchange exchange, SaslStep.Success success, ImapSession session,
                                                   AuthenticateRequest request, Responder responder, byte[] data) {
         if (isAbort(exchange, session, data)) {
-            abortActiveContinuation(exchange, session);
+            closeActiveContinuation(exchange, session);
             no(request, responder, HumanReadableText.AUTHENTICATION_FAILED);
             responder.flush();
             return;
@@ -348,7 +342,7 @@ private void handleTerminalStep(SaslExchange exchange, SaslStep step, ImapSessio
         try {
             handleSaslStep(step, session, request, responder, successLog(request));
         } finally {
-            exchange.close();
+            ImapSaslExchangeTracker.forSession(session).closeExchange(exchange);
         }
     }