[improve][proxy] Refactor broker and proxy auth

Author: nodece

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

@@ -51,6 +51,7 @@ public class BinaryAuthSession {
     private final BinaryAuthContext context;
 
     private AuthResult defaultAuthResult;
+    private boolean supportsAuthRefresh;
 
     public BinaryAuthSession(@NonNull BinaryAuthContext context) {
         this.context = context;
@@ -59,6 +60,8 @@ public BinaryAuthSession(@NonNull BinaryAuthContext context) {
     public CompletableFuture<AuthResult> doAuthentication() {
         var connect = context.getCommandConnect();
         try {
+            supportsAuthRefresh = connect.getFeatureFlags().hasSupportsAuthRefresh() && connect.getFeatureFlags()
+                    .isSupportsAuthRefresh();
             var authData = connect.hasAuthData() ? connect.getAuthData() : emptyArray;
             var clientData = AuthData.of(authData);
             // init authentication
@@ -267,6 +270,47 @@ private CompletableFuture<Void> authenticateOriginalData() {
                 }, context.getExecutor());
     }
 
+    public boolean isExpired() {
+        if (originalAuthState != null) {
+            return originalAuthState.isExpired();
+        }
+        return authState.isExpired();
+    }
+
+    public boolean supportsAuthenticationRefresh(){
+        if (originalPrincipal != null && originalAuthState == null) {
+            // This case is only checked when the authState is expired because we've reached a point where
+            // authentication needs to be refreshed, but the protocol does not support it unless the proxy forwards
+            // the originalAuthData.
+            log.info(
+                    "[{}] Cannot revalidate user credential when using proxy and"
+                            + " not forwarding the credentials.",
+                    context.getRemoteAddress());
+            return false;
+        }
+
+        if (!supportsAuthRefresh) {
+            log.warn("[{}] Client doesn't support auth credentials refresh",
+                    context.getRemoteAddress());
+            return false;
+        }
+
+        return true;
+    }
+
+    public AuthResult refreshAuthentication() throws AuthenticationException {
+        if (originalAuthState != null) {
+            return AuthResult.builder()
+                    .authData(originalAuthState.refreshAuthentication())
+                    .authMethod(originalAuthMethod)
+                    .build();
+        }
+        return AuthResult.builder()
+                .authData(authState.refreshAuthentication())
+                .authMethod(authMethod)
+                .build();
+    }
+
     @Builder
     @Getter
     public static class AuthResult {

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java

@@ -680,6 +680,7 @@ public void testAuthChallengePrincipalChangeFails() throws Exception {
         Object responseConnected = getResponse();
         assertTrue(responseConnected instanceof CommandConnected);
         assertEquals(serverCnx.getState(), State.Connected);
+        assertEquals(serverCnx.getAuthRole(), "pass.client");
         assertEquals(serverCnx.getPrincipal(), "pass.client");
         assertTrue(serverCnx.isActive());
 
@@ -1273,7 +1274,7 @@ private class ClientChannel implements Closeable {
                 4),
                 serverCnx);
         public ClientChannel() {
-            serverCnx.setAuthRole("");
+            serverCnx.clearBinaryAuthSession();
         }
         public void close(){
             if (channel != null && channel.isActive()) {
@@ -1681,6 +1682,51 @@ public void testProducerOnNotOwnedTopic() throws Exception {
         channel.finish();
     }
 
+    private PulsarAuthorizationProvider injectAuth() throws Exception {
+        svcConfig.setAuthorizationEnabled(true);
+        AuthorizationService authorizationService =
+                spyWithClassAndConstructorArgs(AuthorizationService.class, svcConfig, pulsar.getPulsarResources());
+        Field providerField = AuthorizationService.class.getDeclaredField("provider");
+        providerField.setAccessible(true);
+        PulsarAuthorizationProvider authorizationProvider =
+                spyWithClassAndConstructorArgs(PulsarAuthorizationProvider.class, svcConfig,
+                        pulsar.getPulsarResources());
+        providerField.set(authorizationService, authorizationProvider);
+        doReturn(authorizationService).when(brokerService).getAuthorizationService();
+        svcConfig.setAuthorizationEnabled(true);
+
+        AuthenticationService authenticationService = mock(AuthenticationService.class);
+        // use a dummy authentication provider
+        AuthenticationProvider authenticationProvider = new AuthenticationProvider() {
+            @Override
+            public void initialize(ServiceConfiguration config) throws IOException {
+
+            }
+
+            @Override
+            public String authenticate(AuthenticationDataSource authData) throws AuthenticationException {
+                return "role";
+            }
+
+            @Override
+            public String getAuthMethodName() {
+                return "dummy";
+            }
+
+            @Override
+            public void close() throws IOException {
+
+            }
+        };
+
+        String authMethodName = authenticationProvider.getAuthMethodName();
+        when(brokerService.getAuthenticationService()).thenReturn(authenticationService);
+        when(authenticationService.getAuthenticationProvider(authMethodName)).thenReturn(authenticationProvider);
+        svcConfig.setAuthenticationEnabled(true);
+
+        return authorizationProvider;
+    }
+
     @Test(timeOut = 30000)
     public void testProducerCommandWithAuthorizationPositive() throws Exception {
         AuthorizationService authorizationService = mock(AuthorizationService.class);
@@ -1709,32 +1755,27 @@ public void testProducerCommandWithAuthorizationPositive() throws Exception {
 
     @Test(timeOut = 30000)
     public void testNonExistentTopic() throws Exception {
-        AuthorizationService authorizationService =
-                spyWithClassAndConstructorArgs(AuthorizationService.class, svcConfig, pulsar.getPulsarResources());
-        doReturn(authorizationService).when(brokerService).getAuthorizationService();
-        svcConfig.setAuthorizationEnabled(true);
-        svcConfig.setAuthorizationEnabled(true);
-        Field providerField = AuthorizationService.class.getDeclaredField("provider");
-        providerField.setAccessible(true);
-        PulsarAuthorizationProvider authorizationProvider =
-                spyWithClassAndConstructorArgs(PulsarAuthorizationProvider.class, svcConfig,
-                        pulsar.getPulsarResources());
-        providerField.set(authorizationService, authorizationProvider);
+        resetChannel();
+        PulsarAuthorizationProvider authorizationProvider = injectAuth();
         doReturn(CompletableFuture.completedFuture(false)).when(authorizationProvider)
                 .isSuperUser(Mockito.anyString(), Mockito.any(), Mockito.any());
 
+        // Connect
+        ByteBuf connectCommand = Commands.newConnect("dummy", "", null);
+        BinaryAuthSession binaryAuthSession =
+                spyBinaryAuthSession(brokerService.getAuthenticationService(), connectCommand.copy(), svcConfig);
+        when(brokerService.getAuthenticationService().createBinaryAuthSession(any())).thenReturn(binaryAuthSession);
+        channel.writeInbound(connectCommand);
+        Object response = getResponse();
+        assertTrue(response instanceof CommandConnected);
+
         // Test producer creation
-        resetChannel();
-        setChannelConnected();
         ByteBuf newProducerCmd = Commands.newProducer(nonExistentTopicName, 1 /* producer id */, 1 /* request id */,
                 "prod-name", Collections.emptyMap(), false);
         channel.writeInbound(newProducerCmd);
         assertTrue(getResponse() instanceof CommandError);
-        channel.finish();
 
         // Test consumer creation
-        resetChannel();
-        setChannelConnected();
         ByteBuf newSubscribeCmd = Commands.newSubscribe(nonExistentTopicName, //
                 successSubName, 1 /* consumer id */, 1 /* request id */, SubType.Exclusive, 0,
                 "test" /* consumer name */, 0);
@@ -1745,17 +1786,9 @@ public void testNonExistentTopic() throws Exception {
 
     @Test(timeOut = 30000)
     public void testClusterAccess() throws Exception {
-        svcConfig.setAuthorizationEnabled(true);
-        AuthorizationService authorizationService =
-                spyWithClassAndConstructorArgs(AuthorizationService.class, svcConfig, pulsar.getPulsarResources());
-        Field providerField = AuthorizationService.class.getDeclaredField("provider");
-        providerField.setAccessible(true);
-        PulsarAuthorizationProvider authorizationProvider =
-                spyWithClassAndConstructorArgs(PulsarAuthorizationProvider.class, svcConfig,
-                        pulsar.getPulsarResources());
-        providerField.set(authorizationService, authorizationProvider);
-        doReturn(authorizationService).when(brokerService).getAuthorizationService();
-        svcConfig.setAuthorizationEnabled(true);
+        resetChannel();
+
+        PulsarAuthorizationProvider authorizationProvider = injectAuth();
         doReturn(CompletableFuture.completedFuture(false)).when(authorizationProvider)
                 .isSuperUser(Mockito.anyString(), Mockito.any(), Mockito.any());
         doReturn(CompletableFuture.completedFuture(false)).when(authorizationProvider)
@@ -1764,15 +1797,20 @@ public void testClusterAccess() throws Exception {
                 .checkPermission(any(TopicName.class), Mockito.anyString(),
                         any(AuthAction.class));
 
-        resetChannel();
-        setChannelConnected();
+        // Connect
+        ByteBuf connectCommand = Commands.newConnect("dummy", "", null);
+        BinaryAuthSession binaryAuthSession =
+                spyBinaryAuthSession(brokerService.getAuthenticationService(), connectCommand.copy(), svcConfig);
+        when(brokerService.getAuthenticationService().createBinaryAuthSession(any())).thenReturn(binaryAuthSession);
+        channel.writeInbound(connectCommand);
+        Object response = getResponse();
+        assertTrue(response instanceof CommandConnected);
+
         ByteBuf clientCommand = Commands.newProducer(successTopicName, 1 /* producer id */, 1 /* request id */,
                 "prod-name", Collections.emptyMap(), false);
         channel.writeInbound(clientCommand);
         assertTrue(getResponse() instanceof CommandProducerSuccess);
 
-        resetChannel();
-        setChannelConnected();
         clientCommand = Commands.newProducer(topicWithNonLocalCluster, 1 /* producer id */, 1 /* request id */,
                 "prod-name", Collections.emptyMap(), false);
         channel.writeInbound(clientCommand);
@@ -1782,22 +1820,21 @@ public void testClusterAccess() throws Exception {
 
     @Test(timeOut = 30000)
     public void testNonExistentTopicSuperUserAccess() throws Exception {
-        AuthorizationService authorizationService =
-                spyWithClassAndConstructorArgs(AuthorizationService.class, svcConfig, pulsar.getPulsarResources());
-        doReturn(authorizationService).when(brokerService).getAuthorizationService();
-        svcConfig.setAuthorizationEnabled(true);
-        Field providerField = AuthorizationService.class.getDeclaredField("provider");
-        providerField.setAccessible(true);
-        PulsarAuthorizationProvider authorizationProvider =
-                spyWithClassAndConstructorArgs(PulsarAuthorizationProvider.class, svcConfig,
-                        pulsar.getPulsarResources());
-        providerField.set(authorizationService, authorizationProvider);
+        resetChannel();
+        PulsarAuthorizationProvider authorizationProvider = injectAuth();
         doReturn(CompletableFuture.completedFuture(true)).when(authorizationProvider)
                 .isSuperUser(Mockito.anyString(), Mockito.any(), Mockito.any());
 
+        // Connect
+        ByteBuf connectCommand = Commands.newConnect("dummy", "", null);
+        BinaryAuthSession binaryAuthSession =
+                spyBinaryAuthSession(brokerService.getAuthenticationService(), connectCommand.copy(), svcConfig);
+        when(brokerService.getAuthenticationService().createBinaryAuthSession(any())).thenReturn(binaryAuthSession);
+        channel.writeInbound(connectCommand);
+        Object response = getResponse();
+        assertTrue(response instanceof CommandConnected);
+
         // Test producer creation
-        resetChannel();
-        setChannelConnected();
         ByteBuf newProducerCmd = Commands.newProducer(nonExistentTopicName, 1 /* producer id */, 1 /* request id */,
                 "prod-name", Collections.emptyMap(), false);
         channel.writeInbound(newProducerCmd);
@@ -1806,11 +1843,8 @@ public void testNonExistentTopicSuperUserAccess() throws Exception {
         PersistentTopic topicRef = (PersistentTopic) brokerService.getTopicReference(nonExistentTopicName).get();
         assertNotNull(topicRef);
         assertEquals(topicRef.getProducers().size(), 1);
-        channel.finish();
 
         // Test consumer creation
-        resetChannel();
-        setChannelConnected();
         ByteBuf newSubscribeCmd = Commands.newSubscribe(nonExistentTopicName, //
                 successSubName, 1 /* consumer id */, 1 /* request id */, SubType.Exclusive, 0,
                 "test" /* consumer name */, 0 /* avoid reseting cursor */);
@@ -2916,7 +2950,6 @@ protected void resetChannel() throws Exception {
             channel.close().get();
         }
         serverCnx = new ServerCnx(pulsar);
-        serverCnx.setAuthRole("");
         channel = new EmbeddedChannel(new LengthFieldBasedFrameDecoder(
                 maxMessageSize,
                 0,

pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ServerCnxTest.java

@@ -43,13 +43,17 @@
 import io.netty.handler.timeout.ReadTimeoutHandler;
 import io.netty.util.CharsetUtil;
 import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import lombok.Getter;
 import lombok.SneakyThrows;
 import org.apache.pulsar.PulsarVersion;
+import org.apache.pulsar.broker.authentication.AuthenticationDataSource;
+import org.apache.pulsar.broker.authentication.AuthenticationState;
+import org.apache.pulsar.broker.authentication.BinaryAuthSession;
 import org.apache.pulsar.client.api.Authentication;
 import org.apache.pulsar.client.api.AuthenticationDataProvider;
 import org.apache.pulsar.client.api.PulsarClientException;
@@ -100,9 +104,23 @@ public DirectProxyHandler(ProxyService service, ProxyConnection proxyConnection)
         this.inboundChannel = proxyConnection.ctx().channel();
         this.proxyConnection = proxyConnection;
         this.inboundChannelRequestsRate = new Rate();
-        this.originalPrincipal = proxyConnection.clientAuthRole;
-        this.clientAuthData = proxyConnection.clientAuthData;
-        this.clientAuthMethod = proxyConnection.clientAuthMethod;
+        BinaryAuthSession binaryAuthSession = proxyConnection.getBinaryAuthSession();
+        if (binaryAuthSession != null) {
+            AuthenticationState originalAuthState = binaryAuthSession.getOriginalAuthState();
+            boolean forwardOriginal =
+                    originalAuthState != null && service.getConfiguration().isForwardAuthorizationCredentials();
+            AuthenticationDataSource authDataSource = forwardOriginal ? binaryAuthSession.getOriginalAuthData() :
+                    binaryAuthSession.getAuthenticationData();
+            clientAuthData = AuthData.of(authDataSource.getCommandData().getBytes(StandardCharsets.UTF_8));
+            clientAuthMethod =
+                    forwardOriginal ? binaryAuthSession.getOriginalAuthMethod() : binaryAuthSession.getAuthMethod();
+            originalPrincipal =
+                    forwardOriginal ? binaryAuthSession.getOriginalPrincipal() : binaryAuthSession.getAuthRole();
+        } else {
+            originalPrincipal = null;
+            clientAuthData = null;
+            clientAuthMethod = null;
+        }
         this.tlsEnabledWithBroker = service.getConfiguration().isTlsEnabledWithBroker();
         this.tlsHostnameVerificationEnabled = service.getConfiguration().isTlsHostnameVerificationEnabled();
         this.onHandshakeCompleteAction = proxyConnection::cancelKeepAliveTask;