Spike: handle password change in the AuthTokenManager

driver/clirr-ignored-differences.xml

@@ -17,6 +17,19 @@
         <className>org/neo4j/driver/summary/ServerInfo</className>
         <differenceType>7012</differenceType>
         <method>java.lang.String agent()</method>
+        <className>org/neo4j/driver/AuthTokenManager</className>
+    </difference>
+
+    <difference>
+        <className>org/neo4j/driver/AuthTokenManager</className>
+        <differenceType>7002</differenceType>
+        <method>void onExpired(org.neo4j.driver.AuthToken)</method>
+    </difference>
+
+    <difference>
+        <className>org/neo4j/driver/AuthTokenManager</className>
+        <differenceType>7012</differenceType>
+        <method> boolean handleSecurityException(org.neo4j.driver.AuthToken, org.neo4j.driver.exceptions.SecurityException)</method>
     </difference>
 
     <difference>

driver/src/main/java/org/neo4j/driver/AuthTokenManager.java

@@ -19,6 +19,7 @@
 package org.neo4j.driver;
 
 import java.util.concurrent.CompletionStage;
+import org.neo4j.driver.exceptions.SecurityException;
 import org.neo4j.driver.util.Preview;
 
 /**
@@ -56,9 +57,11 @@ public interface AuthTokenManager {
     /**
      * Handles an error notification emitted by the server if the token is expired.
      * <p>
-     * This will be called when driver emits the {@link org.neo4j.driver.exceptions.TokenExpiredRetryableException}.
+     * This will be called when driver emits the {@link org.neo4j.driver.exceptions.AuthenticationException}.
      *
-     * @param authToken the expired token
+     * @param authToken the token
+     * @param exception the security exception
+     * @return handled by the manager.
      */
-    void onExpired(AuthToken authToken);
+    boolean handleSecurityException(AuthToken authToken, SecurityException exception);
 }

driver/src/main/java/org/neo4j/driver/AuthTokenManagers.java

@@ -24,6 +24,7 @@
 import java.util.concurrent.ForkJoinPool;
 import java.util.function.Supplier;
 import org.neo4j.driver.internal.security.ExpirationBasedAuthTokenManager;
+import org.neo4j.driver.internal.security.PasswordChangesAuthTokenManager;
 import org.neo4j.driver.util.Preview;
 
 /**
@@ -42,7 +43,7 @@ private AuthTokenManagers() {}
      * following conditions:
      * <ol>
      *     <li>token's UTC timestamp is expired</li>
-     *     <li>server rejects the current token (see {@link AuthTokenManager#onExpired(AuthToken)})</li>
+     *     <li>server rejects the current token (see {@link AuthTokenManager#handleSecurityException(AuthToken, org.neo4j.driver.exceptions.SecurityException)})</li>
      * </ol>
      * <p>
      * The supplier will be called by a task running in the {@link ForkJoinPool#commonPool()} as documented in the
@@ -62,7 +63,7 @@ public static AuthTokenManager expirationBased(Supplier<AuthTokenAndExpiration>
      * following conditions:
      * <ol>
      *     <li>token's UTC timestamp is expired</li>
-     *     <li>server rejects the current token (see {@link AuthTokenManager#onExpired(AuthToken)})</li>
+     *     <li>server rejects the current token (see {@link AuthTokenManager#handleSecurityException(AuthToken, org.neo4j.driver.exceptions.SecurityException)})</li>
      * </ol>
      * <p>
      * The provided supplier and its completion stages must be non-blocking as documented in the {@link AuthTokenManager}.
@@ -74,4 +75,22 @@ public static AuthTokenManager expirationBasedAsync(
             Supplier<CompletionStage<AuthTokenAndExpiration>> newTokenStageSupplier) {
         return new ExpirationBasedAuthTokenManager(newTokenStageSupplier, Clock.systemUTC());
     }
+
+    /**
+     * Dummy comment
+     * @param newTokenSupplier dummy comment
+     * @return dummy comment
+     */
+    public static AuthTokenManager passwordChanges(Supplier<AuthToken> newTokenSupplier) {
+        return passwordChangesAsync(() -> CompletableFuture.supplyAsync(newTokenSupplier));
+    }
+
+    /**
+     * dummy comment
+     * @param newTokenStageSupplier dummy comment
+     * @return dummy comment
+     */
+    public static AuthTokenManager passwordChangesAsync(Supplier<CompletionStage<AuthToken>> newTokenStageSupplier) {
+        return new PasswordChangesAuthTokenManager(newTokenStageSupplier);
+    }
 }

driver/src/main/java/org/neo4j/driver/internal/async/inbound/InboundMessageDispatcher.java

@@ -36,13 +36,13 @@
 import org.neo4j.driver.Value;
 import org.neo4j.driver.exceptions.AuthorizationExpiredException;
 import org.neo4j.driver.exceptions.ClientException;
+import org.neo4j.driver.exceptions.SecurityException;
 import org.neo4j.driver.exceptions.TokenExpiredException;
 import org.neo4j.driver.exceptions.TokenExpiredRetryableException;
 import org.neo4j.driver.internal.handlers.ResetResponseHandler;
 import org.neo4j.driver.internal.logging.ChannelActivityLogger;
 import org.neo4j.driver.internal.logging.ChannelErrorLogger;
 import org.neo4j.driver.internal.messaging.ResponseMessageHandler;
-import org.neo4j.driver.internal.security.StaticAuthTokenManager;
 import org.neo4j.driver.internal.spi.ResponseHandler;
 import org.neo4j.driver.internal.util.ErrorUtil;
 
@@ -112,6 +112,17 @@ public void handleFailureMessage(String code, String message) {
 
         currentError = ErrorUtil.newNeo4jError(code, message);
 
+        if (currentError instanceof SecurityException securityException) {
+            var authContext = authContext(channel);
+            var authTokenProvider = authContext.getAuthTokenManager();
+            var authToken = authContext.getAuthToken();
+            if (authToken != null && authContext.isManaged()) {
+                if (authTokenProvider.handleSecurityException(authToken, securityException)) {
+                    currentError = toRetryableVersion(securityException);
+                }
+            }
+        }
+
         if (ErrorUtil.isFatal(currentError)) {
             // we should not continue using channel after a fatal error
             // fire error event back to the pipeline and avoid sending RESET
@@ -122,18 +133,7 @@ public void handleFailureMessage(String code, String message) {
         var currentError = this.currentError;
         if (currentError instanceof AuthorizationExpiredException authorizationExpiredException) {
             authorizationStateListener(channel).onExpired(authorizationExpiredException, channel);
-        } else if (currentError instanceof TokenExpiredException tokenExpiredException) {
-            var authContext = authContext(channel);
-            var authTokenProvider = authContext.getAuthTokenManager();
-            if (!(authTokenProvider instanceof StaticAuthTokenManager)) {
-                currentError = new TokenExpiredRetryableException(
-                        tokenExpiredException.code(), tokenExpiredException.getMessage());
-            }
-            var authToken = authContext.getAuthToken();
-            if (authToken != null && authContext.isManaged()) {
-                authTokenProvider.onExpired(authToken);
-            }
-        } else {
+        } else if (!(currentError instanceof TokenExpiredException tokenExpiredException)) {
             // write a RESET to "acknowledge" the failure
             enqueue(new ResetResponseHandler(this));
             channel.writeAndFlush(RESET, channel.voidPromise());
@@ -144,6 +144,10 @@ public void handleFailureMessage(String code, String message) {
         handler.onFailure(currentError);
     }
 
+    private static TokenExpiredRetryableException toRetryableVersion(SecurityException securityException) {
+        return new TokenExpiredRetryableException(securityException.code(), securityException.getMessage());
+    }
+
     @Override
     public void handleIgnoredMessage() {
         log.debug("S: IGNORED");

driver/src/main/java/org/neo4j/driver/internal/security/ExpirationBasedAuthTokenManager.java

@@ -32,6 +32,8 @@
 import org.neo4j.driver.AuthToken;
 import org.neo4j.driver.AuthTokenAndExpiration;
 import org.neo4j.driver.AuthTokenManager;
+import org.neo4j.driver.exceptions.SecurityException;
+import org.neo4j.driver.exceptions.TokenExpiredException;
 
 public class ExpirationBasedAuthTokenManager implements AuthTokenManager {
     private final ReadWriteLock lock = new ReentrantReadWriteLock();
@@ -65,12 +67,16 @@ public CompletionStage<AuthToken> getToken() {
         return validTokenFuture;
     }
 
-    public void onExpired(AuthToken authToken) {
-        executeWithLock(lock.writeLock(), () -> {
-            if (token != null && token.authToken().equals(authToken)) {
-                unsetTokenState();
-            }
-        });
+    public boolean handleSecurityException(AuthToken authToken, SecurityException exception) {
+        if (exception instanceof TokenExpiredException) {
+            executeWithLock(lock.writeLock(), () -> {
+                if (token != null && token.authToken().equals(authToken)) {
+                    unsetTokenState();
+                }
+            });
+            return true;
+        }
+        return false;
     }
 
     private void handleUpstreamResult(AuthTokenAndExpiration authTokenAndExpiration, Throwable throwable) {

driver/src/main/java/org/neo4j/driver/internal/security/PasswordChangesAuthTokenManager.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * 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.
+ */
+package org.neo4j.driver.internal.security;
+
+import static java.util.Objects.requireNonNull;
+import static org.neo4j.driver.internal.util.Futures.failedFuture;
+import static org.neo4j.driver.internal.util.LockUtil.executeWithLock;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Supplier;
+import org.neo4j.driver.AuthToken;
+import org.neo4j.driver.AuthTokenManager;
+import org.neo4j.driver.exceptions.SecurityException;
+import org.neo4j.driver.util.Preview;
+
+@Preview(name = "Password rotation and session auth support")
+public class PasswordChangesAuthTokenManager implements AuthTokenManager {
+    private final ReadWriteLock lock = new ReentrantReadWriteLock();
+    private final Supplier<CompletionStage<AuthToken>> freshTokenSupplier;
+
+    private CompletableFuture<AuthToken> tokenFuture;
+    private AuthToken token;
+
+    public PasswordChangesAuthTokenManager(Supplier<CompletionStage<AuthToken>> freshTokenSupplier) {
+        this.freshTokenSupplier = freshTokenSupplier;
+    }
+
+    @Override
+    public CompletionStage<AuthToken> getToken() {
+        var validTokenFuture = executeWithLock(lock.readLock(), this::getValidTokenFuture);
+        if (validTokenFuture == null) {
+            var fetchFromUpstream = new AtomicBoolean();
+            validTokenFuture = executeWithLock(lock.writeLock(), () -> {
+                if (getValidTokenFuture() == null) {
+                    tokenFuture = new CompletableFuture<>();
+                    token = null;
+                    fetchFromUpstream.set(true);
+                }
+                return tokenFuture;
+            });
+            if (fetchFromUpstream.get()) {
+                getFromUpstream().whenComplete(this::handleUpstreamResult);
+            }
+        }
+        return validTokenFuture;
+    }
+
+    private CompletableFuture<AuthToken> getValidTokenFuture() {
+        CompletableFuture<AuthToken> validTokenFuture = null;
+        if (tokenFuture != null) {
+            validTokenFuture = tokenFuture;
+        }
+        return validTokenFuture;
+    }
+
+    private CompletionStage<AuthToken> getFromUpstream() {
+        CompletionStage<AuthToken> upstreamStage;
+        try {
+            upstreamStage = freshTokenSupplier.get();
+            requireNonNull(upstreamStage, "upstream supplied a null value");
+        } catch (Throwable t) {
+            upstreamStage = failedFuture(t);
+        }
+        return upstreamStage;
+    }
+
+    private void handleUpstreamResult(AuthToken authToken, Throwable throwable) {
+        if (throwable != null) {
+            var previousTokenFuture = executeWithLock(lock.writeLock(), this::unsetTokenState);
+            // notify downstream consumers of the failure
+            previousTokenFuture.completeExceptionally(throwable);
+        } else {
+            var currentTokenFuture = executeWithLock(lock.writeLock(), () -> {
+                token = authToken;
+                return tokenFuture;
+            });
+            currentTokenFuture.complete(authToken);
+        }
+    }
+
+    private CompletableFuture<AuthToken> unsetTokenState() {
+        var previousTokenFuture = tokenFuture;
+        tokenFuture = null;
+        token = null;
+        return previousTokenFuture;
+    }
+
+    @Override
+    public boolean handleSecurityException(AuthToken authToken, SecurityException exception) {
+        executeWithLock(lock.writeLock(), () -> {
+            if (token != null && token.equals(authToken)) {
+                unsetTokenState();
+            }
+        });
+        return true;
+    }
+}

driver/src/main/java/org/neo4j/driver/internal/security/StaticAuthTokenManager.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.neo4j.driver.AuthToken;
 import org.neo4j.driver.AuthTokenManager;
+import org.neo4j.driver.exceptions.SecurityException;
 import org.neo4j.driver.exceptions.TokenExpiredException;
 
 public class StaticAuthTokenManager implements AuthTokenManager {
@@ -44,9 +45,10 @@ public CompletionStage<AuthToken> getToken() {
     }
 
     @Override
-    public void onExpired(AuthToken authToken) {
+    public boolean handleSecurityException(AuthToken authToken, SecurityException exception) {
         if (authToken.equals(this.authToken)) {
             expired.set(true);
         }
+        return false;
     }
 }

driver/src/main/java/org/neo4j/driver/internal/security/ValidatingAuthTokenManager.java

@@ -29,6 +29,7 @@
 import org.neo4j.driver.Logger;
 import org.neo4j.driver.Logging;
 import org.neo4j.driver.exceptions.AuthTokenManagerExecutionException;
+import org.neo4j.driver.exceptions.SecurityException;
 
 public class ValidatingAuthTokenManager implements AuthTokenManager {
     private final Logger log;
@@ -68,19 +69,21 @@ public CompletionStage<AuthToken> getToken() {
     }
 
     @Override
-    public void onExpired(AuthToken authToken) {
+    public boolean handleSecurityException(AuthToken authToken, SecurityException exception) {
         requireNonNull(authToken, "authToken must not be null");
+        requireNonNull(exception, "exception must not be null");
         try {
-            delegate.onExpired(authToken);
+            return delegate.handleSecurityException(authToken, exception);
         } catch (Throwable throwable) {
             log.warn(String.format(
-                    "%s has been thrown by %s.onExpired method",
+                    "%s has been thrown by %s.handleSecurityException method",
                     throwable.getClass().getName(), delegate.getClass().getName()));
             log.debug(
                     String.format(
-                            "%s has been thrown by %s.onExpired method",
+                            "%s has been thrown by %s.handleSecurityException method",
                             throwable.getClass().getName(), delegate.getClass().getName()),
                     throwable);
         }
+        return false;
     }
 }