Update AuthToken rotation to support more auth types

driver/clirr-ignored-differences.xml

@@ -568,4 +568,33 @@
         <method>org.neo4j.driver.AuthTokenAndExpiration expiringAt(long)</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>
+        <className>org/neo4j/driver/AuthTokenManager</className>
+        <differenceType>7002</differenceType>
+        <method>void onExpired(org.neo4j.driver.AuthToken)</method>
+    </difference>
+
+    <difference>
+        <className>org/neo4j/driver/AuthTokenManagers</className>
+        <differenceType>7002</differenceType>
+        <method>org.neo4j.driver.AuthTokenManager expirationBased(java.util.function.Supplier)</method>
+    </difference>
+
+    <difference>
+        <className>org/neo4j/driver/AuthTokenManagers</className>
+        <differenceType>7002</differenceType>
+        <method>org.neo4j.driver.AuthTokenManager expirationBasedAsync(java.util.function.Supplier)</method>
+    </difference>
+
+    <difference>
+        <className>org/neo4j/driver/exceptions/TokenExpiredRetryableException</className>
+        <differenceType>8001</differenceType>
+    </difference>
+
 </differences>

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

@@ -36,14 +36,14 @@ public sealed interface AuthToken permits InternalAuthToken {
     /**
      * Returns a new instance of a type holding both the token and its UTC expiration timestamp.
      * <p>
-     * This is used by the expiration-based implementation of the {@link AuthTokenManager} supplied by the
+     * This is used by the bearer token implementation of the {@link AuthTokenManager} supplied by the
      * {@link AuthTokenManagers}.
      *
      * @param utcExpirationTimestamp the UTC expiration timestamp
      * @return a new instance of a type holding both the token and its UTC expiration timestamp
      * @since 5.8
-     * @see AuthTokenManagers#expirationBased(Supplier)
-     * @see AuthTokenManagers#expirationBasedAsync(Supplier)
+     * @see AuthTokenManagers#bearer(Supplier)
+     * @see AuthTokenManagers#bearerAsync(Supplier)
      */
     @Preview(name = "AuthToken rotation and session auth support")
     default AuthTokenAndExpiration expiringAt(long utcExpirationTimestamp) {

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

@@ -26,12 +26,12 @@
  * A container used by the expiration based {@link AuthTokenManager} implementation provided by the driver, it contains an
  * {@link AuthToken} and its UTC expiration timestamp.
  * <p>
- * This is used by the expiration-based implementation of the {@link AuthTokenManager} supplied by the
+ * This is used by the bearer token implementation of the {@link AuthTokenManager} supplied by the
  * {@link AuthTokenManagers}.
  *
  * @since 5.8
- * @see AuthTokenManagers#expirationBased(Supplier)
- * @see AuthTokenManagers#expirationBasedAsync(Supplier)
+ * @see AuthTokenManagers#bearer(Supplier)
+ * @see AuthTokenManagers#bearerAsync(Supplier)
  */
 @Preview(name = "AuthToken rotation and session auth support")
 public sealed interface AuthTokenAndExpiration permits InternalAuthTokenAndExpiration {

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

@@ -19,6 +19,9 @@
 package org.neo4j.driver;
 
 import java.util.concurrent.CompletionStage;
+import org.neo4j.driver.exceptions.AuthTokenManagerExecutionException;
+import org.neo4j.driver.exceptions.SecurityException;
+import org.neo4j.driver.exceptions.SecurityRetryableException;
 import org.neo4j.driver.util.Preview;
 
 /**
@@ -49,16 +52,27 @@ public interface AuthTokenManager {
      * <p>
      * Failures will surface via the driver API, like {@link Session#beginTransaction()} method and others.
      * @return a stage for a valid token, must not be {@code null} or complete with {@code null}
-     * @see org.neo4j.driver.exceptions.AuthTokenManagerExecutionException
+     * @see AuthTokenManagerExecutionException
      */
     CompletionStage<AuthToken> getToken();
 
     /**
-     * Handles an error notification emitted by the server if the token is expired.
+     * Handles {@link SecurityException} that is created based on the server's security error response by determining if
+     * the given error may be resolved upon next {@link AuthTokenManager#getToken()} invokation.
      * <p>
-     * This will be called when driver emits the {@link org.neo4j.driver.exceptions.TokenExpiredRetryableException}.
+     * If this method returns {@code true}, the driver wraps the original {@link SecurityException} in
+     * {@link SecurityRetryableException}. The managed transaction API (like
+     * {@link Session#executeRead(TransactionCallback)}, etc.) automatically retries its unit of work if no other
+     * condition is violated, while the other query execution APIs surface this error for external handling.
+     * <p>
+     * If this method returns {@code false}, the original error remains unchanged.
+     * <p>
+     * This method must not throw exceptions.
      *
-     * @param authToken the expired token
+     * @param authToken the token
+     * @param exception the security exception
+     * @return {@code true} if the exception should be marked as retryable or {@code false} if it should remain unchanged
+     * @since 5.12
      */
-    void onExpired(AuthToken authToken);
+    boolean handleSecurityException(AuthToken authToken, SecurityException exception);
 }

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

@@ -18,11 +18,17 @@
  */
 package org.neo4j.driver;
 
+import static java.util.Objects.requireNonNull;
+
 import java.time.Clock;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ForkJoinPool;
 import java.util.function.Supplier;
+import org.neo4j.driver.exceptions.AuthenticationException;
+import org.neo4j.driver.exceptions.SecurityException;
+import org.neo4j.driver.exceptions.TokenExpiredException;
 import org.neo4j.driver.internal.security.ExpirationBasedAuthTokenManager;
 import org.neo4j.driver.util.Preview;
 
@@ -36,42 +42,94 @@ public final class AuthTokenManagers {
     private AuthTokenManagers() {}
 
     /**
-     * Returns an {@link AuthTokenManager} that manages {@link AuthToken} instances with UTC expiration timestamp.
+     * Returns an {@link AuthTokenManager} that manages basic {@link AuthToken} instances.
+     * <p>
+     * The implementation will only use the token supplier when it needs a new token instance, which would happen if
+     * the server rejects the current token with {@link AuthenticationException} (see
+     * {@link AuthTokenManager#handleSecurityException(AuthToken, SecurityException)}).
+     * The provided supplier and its completion stages must be non-blocking as documented in the
+     * {@link AuthTokenManager}.
+     *
+     * @param newTokenSupplier a new token stage supplier
+     * @return a new token manager
+     * @since 5.12
+     */
+    public static AuthTokenManager basic(Supplier<AuthToken> newTokenSupplier) {
+        requireNonNull(newTokenSupplier, "newTokenSupplier must not be null");
+        return basicAsync(() -> CompletableFuture.supplyAsync(newTokenSupplier));
+    }
+
+    /**
+     * Returns an {@link AuthTokenManager} that manages basic {@link AuthToken} instances.
+     * <p>
+     * The implementation will only use the token supplier when it needs a new token instance, which would happen if
+     * the server rejects the current token with {@link AuthenticationException} (see
+     * {@link AuthTokenManager#handleSecurityException(AuthToken, SecurityException)}).
+     * The provided supplier and its completion stages must be non-blocking as documented in the
+     * {@link AuthTokenManager}.
+     *
+     * @param newTokenStageSupplier a new token stage supplier
+     * @return a new token manager
+     * @since 5.12
+     */
+    public static AuthTokenManager basicAsync(Supplier<CompletionStage<AuthToken>> newTokenStageSupplier) {
+        requireNonNull(newTokenStageSupplier, "newTokenStageSupplier must not be null");
+        return new ExpirationBasedAuthTokenManager(
+                () -> newTokenStageSupplier.get().thenApply(authToken -> authToken.expiringAt(Long.MAX_VALUE)),
+                Set.of(AuthenticationException.class),
+                Clock.systemUTC());
+    }
+
+    /**
+     * Returns an {@link AuthTokenManager} that manages bearer {@link AuthToken} instances with UTC expiration
+     * timestamp.
      * <p>
      * The implementation will only use the token supplier when it needs a new token instance. This includes the
      * 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 with either {@link TokenExpiredException} or
+     *     {@link AuthenticationException} (see
+     *     {@link AuthTokenManager#handleSecurityException(AuthToken, SecurityException)})</li>
      * </ol>
      * <p>
      * The supplier will be called by a task running in the {@link ForkJoinPool#commonPool()} as documented in the
      * {@link CompletableFuture#supplyAsync(Supplier)}.
      *
      * @param newTokenSupplier a new token supplier
      * @return a new token manager
+     * @since 5.12
      */
-    public static AuthTokenManager expirationBased(Supplier<AuthTokenAndExpiration> newTokenSupplier) {
-        return expirationBasedAsync(() -> CompletableFuture.supplyAsync(newTokenSupplier));
+    public static AuthTokenManager bearer(Supplier<AuthTokenAndExpiration> newTokenSupplier) {
+        requireNonNull(newTokenSupplier, "newTokenSupplier must not be null");
+        return bearerAsync(() -> CompletableFuture.supplyAsync(newTokenSupplier));
     }
 
     /**
-     * Returns an {@link AuthTokenManager} that manages {@link AuthToken} instances with UTC expiration timestamp.
+     * Returns an {@link AuthTokenManager} that manages bearer {@link AuthToken} instances with UTC expiration
+     * timestamp.
      * <p>
      * The implementation will only use the token supplier when it needs a new token instance. This includes the
      * 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 with either {@link TokenExpiredException} or
+     *     {@link AuthenticationException} (see
+     *     {@link AuthTokenManager#handleSecurityException(AuthToken, SecurityException)})</li>
      * </ol>
      * <p>
      * The provided supplier and its completion stages must be non-blocking as documented in the {@link AuthTokenManager}.
      *
      * @param newTokenStageSupplier a new token stage supplier
      * @return a new token manager
+     * @since 5.12
      */
-    public static AuthTokenManager expirationBasedAsync(
+    public static AuthTokenManager bearerAsync(
             Supplier<CompletionStage<AuthTokenAndExpiration>> newTokenStageSupplier) {
-        return new ExpirationBasedAuthTokenManager(newTokenStageSupplier, Clock.systemUTC());
+        requireNonNull(newTokenStageSupplier, "newTokenStageSupplier must not be null");
+        return new ExpirationBasedAuthTokenManager(
+                newTokenStageSupplier,
+                Set.of(TokenExpiredException.class, AuthenticationException.class),
+                Clock.systemUTC());
     }
 }

driver/src/main/java/org/neo4j/driver/exceptions/RetryableException.java

@@ -22,5 +22,6 @@
  * A marker interface for retryable exceptions.
  * <p>
  * This indicates whether an operation that resulted in retryable exception is worth retrying.
+ * @since 5.0
  */
 public interface RetryableException {}

driver/src/main/java/org/neo4j/driver/exceptions/SecurityRetryableException.java

@@ -0,0 +1,76 @@
+/*
+ * 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.exceptions;
+
+import java.io.Serial;
+import java.util.Objects;
+import org.neo4j.driver.AuthToken;
+import org.neo4j.driver.util.Experimental;
+import org.neo4j.driver.util.Preview;
+
+/**
+ * Indicates that the contained {@link SecurityException} is a {@link RetryableException}, which is determined by the
+ * {@link org.neo4j.driver.AuthTokenManager#handleSecurityException(AuthToken, SecurityException)} method.
+ * <p>
+ * The original {@link java.lang.SecurityException} is always available as a {@link Throwable#getCause()}. The
+ * {@link SecurityRetryableException#code()} and {@link SecurityRetryableException#getMessage()} supply the values from
+ * the original exception.
+ *
+ * @since 5.12
+ */
+@Preview(name = "AuthToken rotation and session auth support")
+public class SecurityRetryableException extends SecurityException implements RetryableException {
+    @Serial
+    private static final long serialVersionUID = 3914900631374208080L;
+
+    /**
+     * The original security exception.
+     */
+    private final SecurityException exception;
+
+    /**
+     * Creates a new instance.
+     *
+     * @param exception the original security exception, must not be {@code null}
+     */
+    public SecurityRetryableException(SecurityException exception) {
+        super(exception.getMessage(), exception);
+        this.exception = Objects.requireNonNull(exception);
+    }
+
+    @Override
+    public String code() {
+        return exception.code();
+    }
+
+    @Override
+    public String getMessage() {
+        return exception.getMessage();
+    }
+
+    /**
+     * Returns the original security exception.
+     *
+     * @return the original security exception
+     */
+    @Experimental
+    public SecurityException securityException() {
+        return exception;
+    }
+}