[server] retry authentication retriable when authentication exception occurs.

Author: loserwang1024

fluss-common/src/main/java/com/alibaba/fluss/exception/RetriableAuthenticationException.java

@@ -0,0 +1,28 @@
+/*
+ *  Copyright (c) 2025 Alibaba Group Holding Ltd.
+ *
+ *  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 com.alibaba.fluss.exception;
+
+/**
+ * This exception is thrown if authentication fails with a retriable error.
+ *
+ * @since 0.7
+ */
+public class RetriableAuthenticationException extends AuthenticationException {
+    public RetriableAuthenticationException(String message) {
+        super(message);
+    }
+}

fluss-common/src/main/java/com/alibaba/fluss/security/auth/ClientAuthenticator.java

@@ -18,6 +18,7 @@
 
 import com.alibaba.fluss.annotation.PublicEvolving;
 import com.alibaba.fluss.exception.AuthenticationException;
+import com.alibaba.fluss.shaded.netty4.io.netty.channel.Channel;
 
 import javax.annotation.Nullable;
 
@@ -28,6 +29,9 @@ public interface ClientAuthenticator {
     /** The protocol name of the authenticator, which will send in the AuthenticateRequest. */
     String protocol();
 
+    /** Initialize the authenticator. */
+    default void initialize(AuthenticateContext context) {}
+
     /**
      * * Generates the initial token or calculates a token based on the server's challenge, then
      * sends it back to the server. This method sets the client authentication status as complete if
@@ -75,4 +79,9 @@ public interface ClientAuthenticator {
 
     /** Checks if the authentication from client side is completed. */
     boolean isCompleted();
+
+    /** The context of the authentication process. */
+    interface AuthenticateContext {
+        Channel channel();
+    }
 }

fluss-common/src/main/java/com/alibaba/fluss/security/auth/ServerAuthenticator.java

@@ -19,6 +19,7 @@
 import com.alibaba.fluss.annotation.PublicEvolving;
 import com.alibaba.fluss.exception.AuthenticationException;
 import com.alibaba.fluss.security.acl.FlussPrincipal;
+import com.alibaba.fluss.shaded.netty4.io.netty.channel.Channel;
 
 /**
  * Authenticator for server side.
@@ -30,6 +31,9 @@ public interface ServerAuthenticator {
 
     String protocol();
 
+    /** Initialize the authenticator. */
+    default void initialize(AuthenticateContext context) {}
+
     /**
      * * Generates the challenge based on the client's token, then sends it back to the client. This
      * method sets the server authentication status as complete if the authentication succeeds.
@@ -79,4 +83,9 @@ public interface ServerAuthenticator {
      * complete).
      */
     FlussPrincipal createPrincipal();
+
+    /** The context of the authentication process. */
+    interface AuthenticateContext {
+        Channel channel();
+    }
 }

fluss-common/src/test/java/com/alibaba/fluss/utils/ExponentialBackoffTest.java

@@ -52,7 +52,7 @@ public void testExponentialBackoffWithoutJitter() {
         ExponentialBackoff exponentialBackoff = new ExponentialBackoff(100, 2, 400, 0.0);
         assertThat(exponentialBackoff.backoff(0)).isEqualTo(100);
         assertThat(exponentialBackoff.backoff(1)).isEqualTo(200);
-        assertThat(exponentialBackoff.backoff(2)).isEqualTo(300);
+        assertThat(exponentialBackoff.backoff(2)).isEqualTo(400);
         assertThat(exponentialBackoff.backoff(3)).isEqualTo(400);
     }
 }

fluss-rpc/src/main/java/com/alibaba/fluss/rpc/netty/client/ServerConnection.java

@@ -20,6 +20,7 @@
 import com.alibaba.fluss.exception.DisconnectException;
 import com.alibaba.fluss.exception.FlussRuntimeException;
 import com.alibaba.fluss.exception.NetworkException;
+import com.alibaba.fluss.exception.RetriableAuthenticationException;
 import com.alibaba.fluss.rpc.messages.ApiMessage;
 import com.alibaba.fluss.rpc.messages.ApiVersionsRequest;
 import com.alibaba.fluss.rpc.messages.ApiVersionsResponse;
@@ -38,6 +39,7 @@
 import com.alibaba.fluss.shaded.netty4.io.netty.channel.Channel;
 import com.alibaba.fluss.shaded.netty4.io.netty.channel.ChannelFuture;
 import com.alibaba.fluss.shaded.netty4.io.netty.channel.ChannelFutureListener;
+import com.alibaba.fluss.utils.ExponentialBackoff;
 import com.alibaba.fluss.utils.MapUtils;
 
 import org.slf4j.Logger;
@@ -52,6 +54,7 @@
 import java.util.ArrayDeque;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 /** Connection to a Netty server used by the {@link NettyClient}. */
@@ -66,6 +69,7 @@ final class ServerConnection {
     private final CompletableFuture<Void> closeFuture = new CompletableFuture<>();
     private final ConnectionMetricGroup connectionMetricGroup;
     private final ClientAuthenticator authenticator;
+    private final ExponentialBackoff backoff;
 
     private final Object lock = new Object();
 
@@ -87,6 +91,9 @@ final class ServerConnection {
     @GuardedBy("lock")
     private ServerApiVersions serverApiVersions;
 
+    @GuardedBy("lock")
+    private int retryAuthCount = 0;
+
     ServerConnection(
             Bootstrap bootstrap,
             ServerNode node,
@@ -99,6 +106,7 @@ final class ServerConnection {
                 .connect(node.host(), node.port())
                 .addListener((ChannelFutureListener) this::establishConnection);
         this.authenticator = authenticator;
+        this.backoff = new ExponentialBackoff(100L, 2, 5000L, 0.2);
     }
 
     public ServerNode getServerNode() {
@@ -366,13 +374,18 @@ private void handleApiVersionsResponse(ApiMessage response, Throwable cause) {
         synchronized (lock) {
             serverApiVersions =
                     new ServerApiVersions(((ApiVersionsResponse) response).getApiVersionsList());
-            LOG.debug("Begin to authenticate with protocol {}", authenticator.protocol());
             // send initial token
-            sendAuthenticate(new byte[0]);
+            sendInitialToken();
         }
     }
 
-    private void sendAuthenticate(byte[] challenge) {
+    private void sendInitialToken() {
+        authenticator.initialize(new DefaultAuthenticateContext());
+        LOG.debug("Begin to authenticate with protocol {}", authenticator.protocol());
+        sendAuthenticateRequest(new byte[0]);
+    }
+
+    private void sendAuthenticateRequest(byte[] challenge) {
         try {
             if (!authenticator.isCompleted()) {
                 byte[] token = authenticator.authenticate(challenge);
@@ -404,7 +417,16 @@ private void sendAuthenticate(byte[] challenge) {
 
     private void handleAuthenticateResponse(ApiMessage response, Throwable cause) {
         if (cause != null) {
-            close(cause);
+            if (cause instanceof RetriableAuthenticationException) {
+                LOG.warn("Authentication failed, retrying {} times", retryAuthCount, cause);
+                channel.eventLoop()
+                        .schedule(
+                                this::sendInitialToken,
+                                backoff.backoff(retryAuthCount++),
+                                TimeUnit.MILLISECONDS);
+            } else {
+                close(cause);
+            }
             return;
         }
         if (!(response instanceof AuthenticateResponse)) {
@@ -415,7 +437,7 @@ private void handleAuthenticateResponse(ApiMessage response, Throwable cause) {
         synchronized (lock) {
             AuthenticateResponse authenticateResponse = (AuthenticateResponse) response;
             if (authenticateResponse.hasChallenge()) {
-                sendAuthenticate(((AuthenticateResponse) response).getChallenge());
+                sendAuthenticateRequest(((AuthenticateResponse) response).getChallenge());
             } else if (authenticator.isCompleted()) {
                 switchState(ConnectionState.READY);
             } else {
@@ -521,4 +543,11 @@ ByteBuf toByteBuf(ByteBufAllocator allocator) {
             return MessageCodec.encodeRequest(allocator, apiKey, apiVersion, requestId, request);
         }
     }
+
+    private class DefaultAuthenticateContext implements ClientAuthenticator.AuthenticateContext {
+        @Override
+        public Channel channel() {
+            return channel;
+        }
+    }
 }

fluss-rpc/src/main/java/com/alibaba/fluss/rpc/netty/server/NettyServerHandler.java

@@ -19,6 +19,7 @@
 import com.alibaba.fluss.annotation.VisibleForTesting;
 import com.alibaba.fluss.exception.AuthenticationException;
 import com.alibaba.fluss.exception.NetworkException;
+import com.alibaba.fluss.exception.RetriableAuthenticationException;
 import com.alibaba.fluss.record.send.Send;
 import com.alibaba.fluss.rpc.messages.ApiMessage;
 import com.alibaba.fluss.rpc.messages.AuthenticateRequest;
@@ -32,6 +33,7 @@
 import com.alibaba.fluss.security.auth.ServerAuthenticator;
 import com.alibaba.fluss.shaded.netty4.io.netty.buffer.ByteBuf;
 import com.alibaba.fluss.shaded.netty4.io.netty.buffer.ByteBufAllocator;
+import com.alibaba.fluss.shaded.netty4.io.netty.channel.Channel;
 import com.alibaba.fluss.shaded.netty4.io.netty.channel.ChannelFutureListener;
 import com.alibaba.fluss.shaded.netty4.io.netty.channel.ChannelHandlerContext;
 import com.alibaba.fluss.shaded.netty4.io.netty.channel.ChannelInboundHandlerAdapter;
@@ -162,6 +164,7 @@ public void channelActive(ChannelHandlerContext ctx) throws Exception {
         super.channelActive(ctx);
         this.ctx = ctx;
         this.remoteAddress = ctx.channel().remoteAddress();
+        authenticator.initialize(new DefaultAuthenticateContext());
         switchState(
                 authenticator.isCompleted()
                         ? ConnectionState.READY
@@ -316,14 +319,27 @@ private void handleAuthenticateRequest(
         }
 
         AuthenticateResponse authenticateResponse = new AuthenticateResponse();
-        if (!authenticator.isCompleted()) {
-            byte[] token = authenticateRequest.getToken();
-            byte[] challenge = authenticator.evaluateResponse(token);
-            if (!authenticator.isCompleted() && challenge != null) {
-                authenticateResponse.setChallenge(challenge);
+        try {
+            if (!authenticator.isCompleted()) {
+                byte[] token = authenticateRequest.getToken();
+                byte[] challenge = authenticator.evaluateResponse(token);
+                if (challenge != null) {
+                    authenticateResponse.setChallenge(challenge);
+                }
+            }
+            future.complete(authenticateResponse);
+        } catch (AuthenticationException e) {
+            if (e instanceof RetriableAuthenticationException) {
+                LOG.warn(
+                        "Authentication from {} failed due to a retriable exception: {}. Reinitializing authenticator for subsequent retries.",
+                        ctx.channel().remoteAddress(),
+                        e.getMessage(),
+                        e);
+                authenticator.initialize(new DefaultAuthenticateContext());
             }
+
+            future.completeExceptionally(e);
         }
-        future.complete(authenticateResponse);
 
         if (authenticator.isCompleted()) {
             switchState(ConnectionState.READY);
@@ -349,4 +365,11 @@ public boolean isAuthenticating() {
             return this == AUTHENTICATING;
         }
     }
+
+    private class DefaultAuthenticateContext implements ServerAuthenticator.AuthenticateContext {
+        @Override
+        public Channel channel() {
+            return ctx.channel();
+        }
+    }
 }

fluss-rpc/src/main/java/com/alibaba/fluss/rpc/protocol/Errors.java

@@ -53,6 +53,7 @@
 import com.alibaba.fluss.exception.PartitionAlreadyExistsException;
 import com.alibaba.fluss.exception.PartitionNotExistException;
 import com.alibaba.fluss.exception.RecordTooLargeException;
+import com.alibaba.fluss.exception.RetriableAuthenticationException;
 import com.alibaba.fluss.exception.SchemaNotExistException;
 import com.alibaba.fluss.exception.SecurityDisabledException;
 import com.alibaba.fluss.exception.SecurityTokenException;
@@ -61,7 +62,6 @@
 import com.alibaba.fluss.exception.TableNotExistException;
 import com.alibaba.fluss.exception.TableNotPartitionedException;
 import com.alibaba.fluss.exception.TimeoutException;
-import com.alibaba.fluss.exception.TooManyBucketsException;
 import com.alibaba.fluss.exception.TooManyPartitionsException;
 import com.alibaba.fluss.exception.UnknownServerException;
 import com.alibaba.fluss.exception.UnknownTableOrBucketException;
@@ -198,8 +198,10 @@ public enum Errors {
     AUTHENTICATE_EXCEPTION(46, "Authentication failed.", AuthenticationException::new),
     SECURITY_DISABLED_EXCEPTION(47, "Security is disabled.", SecurityDisabledException::new),
     AUTHORIZATION_EXCEPTION(48, "Authorization failed", AuthorizationException::new),
-    BUCKET_MAX_NUM_EXCEPTION(
-            49, "Exceed the maximum number of buckets", TooManyBucketsException::new);
+    RETRIABLE_AUTHENTICATE_EXCEPTION(
+            49,
+            "Authentication failed with retriable exception. ",
+            RetriableAuthenticationException::new);
 
     private static final Logger LOG = LoggerFactory.getLogger(Errors.class);
 

fluss-rpc/src/test/java/com/alibaba/fluss/rpc/netty/authenticate/AuthenticationTest.java

@@ -122,6 +122,17 @@ void testNoChallengeBeforeClientComplete() throws Exception {
         }
     }
 
+    @Test
+    void testRetirableAuthenticateException() throws Exception {
+        Configuration clientConfig = new Configuration();
+        clientConfig.set(ConfigOptions.CLIENT_SECURITY_PROTOCOL, "mutual");
+        clientConfig.setString("client.security.mutual.error-type", "RETRIABLE_EXCEPTION");
+        try (NettyClient nettyClient =
+                new NettyClient(clientConfig, TestingClientMetricGroup.newInstance())) {
+            verifyGetTableNamesList(nettyClient, mutualAuthServerNode);
+        }
+    }
+
     @Test
     void testClientLackAuthenticateProtocol() throws Exception {
         Configuration clientConfig = new Configuration();