[improve][broker] AuthenticationState authenticate data once

pulsar-broker-auth-sasl/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderSasl.java

@@ -125,6 +125,12 @@ public void close() throws IOException {
     public AuthenticationState newAuthState(AuthData authData,
                                             SocketAddress remoteAddress,
                                             SSLSession sslSession) throws AuthenticationException {
+        return newAuthState(remoteAddress, sslSession);
+    }
+
+    @Override
+    public AuthenticationState newAuthState(SocketAddress remoteAddress,
+                                            SSLSession sslSession) throws AuthenticationException {
         try {
             PulsarSaslServer server = new PulsarSaslServer(jaasCredentialsContainer.getSubject(), allowedIdsPattern);
             return new SaslAuthenticationState(server);
@@ -269,7 +275,7 @@ public boolean authenticateHttpRequest(HttpServletRequest request, HttpServletRe
             // no role token, do sasl auth
             // need new authState
             if (state == null || request.getHeader(SASL_HEADER_STATE).equalsIgnoreCase(SASL_STATE_CLIENT_INIT)) {
-                state = newAuthState(null, null, null);
+                state = newAuthState(null, null);
                 authStates.put(state.getStateId(), state);
             }
             checkState(request.getHeader(SASL_AUTH_TOKEN) != null,

pulsar-broker-auth-sasl/src/test/java/org/apache/pulsar/broker/authentication/SaslAuthenticateTest.java

@@ -262,7 +262,7 @@ public void testSaslServerAndClientAuth() throws Exception {
                 .getAuthenticationProvider(SaslConstants.AUTH_METHOD_NAME));
         AuthenticationProviderSasl saslServer =
             (AuthenticationProviderSasl) providerList.getProviders().get(0);
-        AuthenticationState authState = saslServer.newAuthState(null, null, null);
+        AuthenticationState authState = saslServer.newAuthState(null, null);
 
         // auth between server and client.
         // first time auth
@@ -288,7 +288,7 @@ public void testSaslServerAndClientAuth() throws Exception {
 
         // another server could not serve old client
         try {
-            AuthenticationState authState2 = saslServer.newAuthState(null, null, null);
+            AuthenticationState authState2 = saslServer.newAuthState(null, null);
             AuthData serverData3 = authState2.authenticate(initData1);
             fail("Expected fail. server is auth old client data");
         } catch (Exception e) {

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProvider.java

@@ -91,9 +91,21 @@ default String authenticate(AuthenticationDataSource authData) throws Authentica
         throw new AuthenticationException("Not supported");
     }
 
+    /**
+     * Create an {@link AuthenticationState} for the given connection information.
+     * @return an {@link AuthenticationState} object without any configured {@link AuthData}.
+     */
+    default AuthenticationState newAuthState(SocketAddress remoteAddress,
+                                             SSLSession sslSession) throws AuthenticationException {
+        return new OneStageAuthenticationState(remoteAddress, sslSession, this);
+    }
+
     /**
      * Create an authentication data State use passed in AuthenticationDataSource.
+     * @deprecated use {@link #newAuthState(SocketAddress, SSLSession)}, and note that the implementation requires
+     * users call {@link AuthenticationState#authenticate(AuthData)}.
      */
+    @Deprecated(since = "2.12.0")
     default AuthenticationState newAuthState(AuthData authData,
                                              SocketAddress remoteAddress,
                                              SSLSession sslSession)

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java

@@ -160,6 +160,12 @@ public String authenticate(AuthenticationDataSource authData) throws Authenticat
         }
     }
 
+    @Override
+    public AuthenticationState newAuthState(SocketAddress remoteAddress, SSLSession sslSession)
+            throws AuthenticationException {
+        return new TokenAuthenticationState(this, remoteAddress, sslSession);
+    }
+
     @Override
     public AuthenticationState newAuthState(AuthData authData, SocketAddress remoteAddress, SSLSession sslSession)
             throws AuthenticationException {
@@ -314,21 +320,45 @@ private String getTokenAudience(ServiceConfiguration conf) throws IllegalArgumen
 
     private static final class TokenAuthenticationState implements AuthenticationState {
         private final AuthenticationProviderToken provider;
+        private final SocketAddress remoteAddress;
+        private final SSLSession sslSession;
         private AuthenticationDataSource authenticationDataSource;
         private Jwt<?, Claims> jwt;
         private long expiration;
 
+        TokenAuthenticationState(
+                AuthenticationProviderToken provider,
+                SocketAddress remoteAddress,
+                SSLSession sslSession) {
+            this.provider = provider;
+            this.remoteAddress = remoteAddress;
+            this.sslSession = sslSession;
+        }
+
+        /**
+         * @deprecated this class no longer stores the {@link AuthData} passed in on initialization, so this constructor
+         * is deprecated. In order to maintain some backwards compatibility, this constructor validates the
+         * parameterized {@link AuthData}.
+         */
+        @Deprecated(since = "2.12.0")
         TokenAuthenticationState(
                 AuthenticationProviderToken provider,
                 AuthData authData,
                 SocketAddress remoteAddress,
                 SSLSession sslSession) throws AuthenticationException {
-            this.provider = provider;
+            this(provider, remoteAddress, sslSession);
             String token = new String(authData.getBytes(), UTF_8);
             this.authenticationDataSource = new AuthenticationDataCommand(token, remoteAddress, sslSession);
             this.checkExpiration(token);
         }
 
+        /**
+         * @deprecated method was only ever used for
+         * {@link AuthenticationProvider#newHttpAuthState(HttpServletRequest)}. That state object wasn't used other than
+         * to retrieve the {@link AuthenticationDataSource}, so this constructor is deprecated now. In order to maintain
+         * some backwards compatibility, this constructor validates the parameterized {@link AuthData}.
+         */
+        @Deprecated(since = "2.12.0")
         TokenAuthenticationState(
                 AuthenticationProviderToken provider,
                 HttpServletRequest request) throws AuthenticationException {
@@ -342,10 +372,17 @@ private static final class TokenAuthenticationState implements AuthenticationSta
             String token = httpHeaderValue.substring(HTTP_HEADER_VALUE_PREFIX.length());
             this.authenticationDataSource = new AuthenticationDataHttps(request);
             this.checkExpiration(token);
+
+            // These are not used when this constructor is invoked, set them to null.
+            this.sslSession = null;
+            this.remoteAddress = null;
         }
 
         @Override
         public String getAuthRole() throws AuthenticationException {
+            if (jwt == null) {
+                throw new AuthenticationException("Must authenticate before calling getAuthRole");
+            }
             return provider.getPrincipal(jwt);
         }
 
@@ -358,6 +395,7 @@ public String getAuthRole() throws AuthenticationException {
         public AuthData authenticate(AuthData authData) throws AuthenticationException {
             String token = new String(authData.getBytes(), UTF_8);
             checkExpiration(token);
+            this.authenticationDataSource = new AuthenticationDataCommand(token, remoteAddress, sslSession);
             return null;
         }
 
@@ -378,8 +416,7 @@ public AuthenticationDataSource getAuthDataSource() {
 
         @Override
         public boolean isComplete() {
-            // The authentication of tokens is always done in one single stage
-            return true;
+            return jwt != null;
         }
 
         @Override

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationState.java

@@ -18,22 +18,31 @@
  */
 package org.apache.pulsar.broker.authentication;
 
+import java.net.SocketAddress;
 import java.util.concurrent.CompletableFuture;
 import javax.naming.AuthenticationException;
+import javax.net.ssl.SSLSession;
 import org.apache.pulsar.common.api.AuthData;
 import org.apache.pulsar.common.util.FutureUtil;
 
 /**
  * Interface for authentication state.
- *
- * It tell broker whether the authentication is completed or not,
- * if completed, what is the AuthRole is.
+ * <p>
+ * Pulsar integrates with this class in the following order:
+ * 1. Initializing the class by calling {@link AuthenticationProvider#newAuthState(SocketAddress, SSLSession)}
+ * 2. Calling {@link #authenticate(AuthData)} in a loop until the result is null.
+ * 3. Calling {@link #getAuthRole()} and {@link #getAuthDataSource()} to use for authentication.
+ * 4. Calling {@link #refreshAuthentication()} when {@link #isExpired()} returns true.
+ * 5. GoTo step 3.
  */
 public interface AuthenticationState {
     /**
      * After the authentication between client and broker completed,
      * get authentication role represent for the client.
      * It should throw exception if auth not complete.
+     *
+     * <p>Users of this interface must call {@link #authenticateAsync(AuthData)} or {@link #authenticate(AuthData)}
+     * before calling this method.</p>
      */
     String getAuthRole() throws AuthenticationException;
 
@@ -71,7 +80,9 @@ default CompletableFuture<AuthData> authenticateAsync(AuthData authData) {
     }
 
     /**
-     * Return AuthenticationDataSource.
+     * <p>Users of this interface must call {@link #authenticateAsync(AuthData)} or {@link #authenticate(AuthData)}
+     * before calling this method.</p>
+     * @return AuthenticationDataSource.
      */
     AuthenticationDataSource getAuthDataSource();
 

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/OneStageAuthenticationState.java

@@ -26,33 +26,74 @@
 import org.apache.pulsar.common.api.AuthData;
 
 /**
- * Interface for authentication state.
- *
- * It tell broker whether the authentication is completed or not,
- * if completed, what is the AuthRole is.
+ * A class to track single stage authentication. This class assumes that:
+ * 1. {@link #authenticate(AuthData)} is called once and then authentication is completed
+ * 2. Authentication does not expire, so {@link #isExpired()} always returns false.
+ * <p>
+ * See {@link AuthenticationState} for Pulsar's contract on how this interface is used by Pulsar.
  */
 public class OneStageAuthenticationState implements AuthenticationState {
 
-    private final AuthenticationDataSource authenticationDataSource;
-    private final String authRole;
+    private AuthenticationDataSource authenticationDataSource;
+    private final SocketAddress remoteAddress;
+    private final SSLSession sslSession;
+    private final AuthenticationProvider provider;
+    private String authRole;
+
+    /**
+     * Constructor for a {@link OneStageAuthenticationState} where there is no authentication performed on
+     * initialization.
+     * @param remoteAddress - remoteAddress associated with the {@link AuthenticationState}
+     * @param sslSession - sslSession associated with the {@link AuthenticationState}
+     * @param provider - {@link AuthenticationProvider} to use to verify {@link AuthData}
+     */
+    public OneStageAuthenticationState(SocketAddress remoteAddress,
+                                       SSLSession sslSession,
+                                       AuthenticationProvider provider) {
+        this.provider = provider;
+        this.remoteAddress = remoteAddress;
+        this.sslSession = sslSession;
+    }
 
+    /**
+     * @deprecated use OneStageAuthenticationState constructor without {@link AuthData}. In order to maintain some
+     * backwards compatibility, this constructor validates the parameterized {@link AuthData}.
+     */
+    @Deprecated(since = "2.12.0")
     public OneStageAuthenticationState(AuthData authData,
                                        SocketAddress remoteAddress,
                                        SSLSession sslSession,
                                        AuthenticationProvider provider) throws AuthenticationException {
         this.authenticationDataSource = new AuthenticationDataCommand(
             new String(authData.getBytes(), UTF_8), remoteAddress, sslSession);
         this.authRole = provider.authenticate(authenticationDataSource);
+        this.provider = provider;
+        this.remoteAddress = remoteAddress;
+        this.sslSession = sslSession;
     }
 
+    /**
+     * @deprecated method was only ever used for
+     * {@link AuthenticationProvider#newHttpAuthState(HttpServletRequest)}. That state object wasn't used other than
+     * to retrieve the {@link AuthenticationDataSource}, so this constructor is deprecated now. In order to maintain
+     * some backwards compatibility, this constructor validates the parameterized {@link AuthData}.
+     */
+    @Deprecated(since = "2.12.0")
     public OneStageAuthenticationState(HttpServletRequest request, AuthenticationProvider provider)
             throws AuthenticationException {
         this.authenticationDataSource = new AuthenticationDataHttps(request);
         this.authRole = provider.authenticate(authenticationDataSource);
+        // These are not used when this constructor is invoked, set them to null.
+        this.provider = null;
+        this.remoteAddress = null;
+        this.sslSession = null;
     }
 
     @Override
-    public String getAuthRole() {
+    public String getAuthRole() throws AuthenticationException {
+        if (authRole == null) {
+            throw new AuthenticationException("Must authenticate before calling getAuthRole");
+        }
         return authRole;
     }
 
@@ -62,12 +103,15 @@ public AuthenticationDataSource getAuthDataSource() {
     }
 
     @Override
-    public AuthData authenticate(AuthData authData) {
+    public AuthData authenticate(AuthData authData) throws AuthenticationException {
+        this.authenticationDataSource = new AuthenticationDataCommand(
+                new String(authData.getBytes(), UTF_8), remoteAddress, sslSession);
+        this.authRole = provider.authenticate(authenticationDataSource);
         return null;
     }
 
     @Override
     public boolean isComplete() {
-        return true;
+        return authRole != null;
     }
 }

pulsar-broker-common/src/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderBasicTest.java

@@ -19,6 +19,7 @@
 package org.apache.pulsar.broker.authentication;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
 import com.google.common.io.Resources;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -39,9 +40,15 @@ public class AuthenticationProviderBasicTest {
     public AuthenticationProviderBasicTest() throws IOException {
     }
 
+    @SuppressWarnings("deprecation")
     private void testAuthenticate(AuthenticationProviderBasic provider) throws AuthenticationException {
         AuthData authData = AuthData.of("superUser2:superpassword".getBytes(StandardCharsets.UTF_8));
+        // Test legacy newAuthState method, which authenticated data
         provider.newAuthState(authData, null, null);
+        // Test new newAuthState method, which requires calling authenticate method to authenticate data
+        AuthenticationState authState = provider.newAuthState(null, null);
+        AuthData challenge = authState.authenticate(authData);
+        assertNull(challenge, "Should not produce a challenge result");
     }
 
     @Test