Async authorization check while creation of producer/consumer (#98)

pulsar-broker-common/src/main/java/com/yahoo/pulsar/broker/authorization/AuthorizationManager.java

@@ -16,8 +16,8 @@
 package com.yahoo.pulsar.broker.authorization;
 
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -26,7 +26,6 @@ import com.yahoo.pulsar.broker.ServiceConfiguration;
 import com.yahoo.pulsar.broker.cache.ConfigurationCacheService;
 import com.yahoo.pulsar.common.naming.DestinationName;
 import com.yahoo.pulsar.common.policies.data.AuthAction;
-import com.yahoo.pulsar.common.policies.data.Policies;
 
 /**
  */
@@ -50,9 +49,19 @@ public class AuthorizationManager {
      * @param role
      *            the app id used to send messages to the destination.
      */
-    public boolean canProduce(DestinationName destination, String role) {
+    public CompletableFuture<Boolean> canProduceAsync(DestinationName destination, String role) {
         return checkAuthorization(destination, role, AuthAction.produce);
     }
+    
+    public boolean canProduce(DestinationName destination, String role) {
+        try {
+            return canProduceAsync(destination, role).get();
+        } catch (Exception e) {
+            log.warn("Producer-client  with Role - {} failed to get permissions for destination - {}", role,
+                    destination, e);
+            return false;
+        }
+    }
 
     /**
      * Check if the specified role has permission to receive messages from the specified fully qualified destination
@@ -63,9 +72,19 @@ public class AuthorizationManager {
      * @param role
      *            the app id used to receive messages from the destination.
      */
-    public boolean canConsume(DestinationName destination, String role) {
+    public CompletableFuture<Boolean> canConsumeAsync(DestinationName destination, String role) {
         return checkAuthorization(destination, role, AuthAction.consume);
     }
+    
+    public boolean canConsume(DestinationName destination, String role) {
+        try {
+            return canConsumeAsync(destination, role).get();
+        } catch (Exception e) {
+            log.warn("Consumer-client  with Role - {} failed to get permissions for destination - {}", role,
+                    destination, e);
+            return false;
+        }
+    }
 
     /**
      * Check whether the specified role can perform a lookup for the specified destination.
@@ -80,10 +99,14 @@ public class AuthorizationManager {
         return canProduce(destination, role) || canConsume(destination, role);
     }
 
-    private boolean checkAuthorization(DestinationName destination, String role, AuthAction action) {
-        if (isSuperUser(role))
-            return true;
-        return checkPermission(destination, role, action) && checkCluster(destination);
+    private CompletableFuture<Boolean> checkAuthorization(DestinationName destination, String role,
+            AuthAction action) {
+        if (isSuperUser(role)) {
+            return CompletableFuture.completedFuture(true);
+        } else {
+            return checkPermission(destination, role, action)
+                    .thenApply(isPermission -> isPermission && checkCluster(destination));
+        }
     }
 
     private boolean checkCluster(DestinationName destination) {
@@ -98,38 +121,49 @@ public class AuthorizationManager {
         }
     }
 
-    public boolean checkPermission(DestinationName destination, String role, AuthAction action) {
+    public CompletableFuture<Boolean> checkPermission(DestinationName destination, String role,
+            AuthAction action) {
+        CompletableFuture<Boolean> permissionFuture = new CompletableFuture<>();
         try {
-            Optional<Policies> policies = configCache.policiesCache().get(POLICY_ROOT + destination.getNamespace());
-            if (!policies.isPresent()) {
-                if (log.isDebugEnabled()) {
-                    log.debug("Policies node couldn't be found for destination : {}", destination);
+            configCache.policiesCache().getAsync(POLICY_ROOT + destination.getNamespace()).thenAccept(policies -> {
+                if (!policies.isPresent()) {
+                    if (log.isDebugEnabled()) {
+                        log.debug("Policies node couldn't be found for destination : {}", destination);
+                    }
+                    permissionFuture.complete(false);
+                } else {
+                    Set<AuthAction> namespaceActions = policies.get().auth_policies.namespace_auth.get(role);
+                    if (namespaceActions != null && namespaceActions.contains(action)) {
+                        // The role has namespace level permission
+                        permissionFuture.complete(true);
+                    } else {
+                        Map<String, Set<AuthAction>> roles = policies.get().auth_policies.destination_auth
+                                .get(destination.toString());
+                        if (roles == null) {
+                            // Destination has no custom policy
+                            permissionFuture.complete(false);
+                        } else {
+                            Set<AuthAction> resourceActions = roles.get(role);
+                            if (resourceActions != null && resourceActions.contains(action)) {
+                                // The role has destination level permission
+                                permissionFuture.complete(true);
+                            } else {
+                                permissionFuture.complete(false);
+                            }
+                        }
+                    }
                 }
-                return false;
-            }
-
-            Set<AuthAction> namespaceActions = policies.get().auth_policies.namespace_auth.get(role);
-            if (namespaceActions != null && namespaceActions.contains(action)) {
-                // The role has namespace level permission
-                return true;
-            }
-
-            Map<String, Set<AuthAction>> roles = policies.get().auth_policies.destination_auth.get(destination.toString());
-            if (roles == null) {
-                // Destination has no custom policy
-                return false;
-            }
-
-            Set<AuthAction> resourceActions = roles.get(role);
-            if (resourceActions != null && resourceActions.contains(action)) {
-                // The role has destination level permission
-                return true;
-            }
-            return false;
+            }).exceptionally(ex -> {
+                log.warn("Client  with Role - {} failed to get permissions for destination - {}", role, destination,
+                        ex);
+                permissionFuture.complete(false);
+                return null;
+            });
         } catch (Exception e) {
             log.warn("Client  with Role - {} failed to get permissions for destination - {}", role, destination, e);
-            return false;
+            permissionFuture.complete(false);
         }
+        return permissionFuture;
     }
 
     /**

pulsar-broker/src/test/java/com/yahoo/pulsar/broker/service/ServerCnxTest.java

@@ -381,7 +381,7 @@ public class ServerCnxTest {
     @Test(timeOut = 30000)
     public void testProducerCommandWithAuthorizationPositive() throws Exception {
         AuthorizationManager authorizationManager = mock(AuthorizationManager.class);
-        doReturn(true).when(authorizationManager).canProduce(Mockito.any(), Mockito.any());
+        doReturn(CompletableFuture.completedFuture(true)).when(authorizationManager).canProduceAsync(Mockito.any(), Mockito.any());
         doReturn(authorizationManager).when(brokerService).getAuthorizationManager();
         doReturn(true).when(brokerService).isAuthenticationEnabled();
         resetChannel();
@@ -409,7 +409,7 @@ public class ServerCnxTest {
         ConfigurationCacheService configCacheService = mock(ConfigurationCacheService.class);
         doReturn(configCacheService).when(pulsar).getConfigurationCache();
         doReturn(zkDataCache).when(configCacheService).policiesCache();
-        doThrow(new NoNodeException()).when(zkDataCache).get(matches(".*nonexistent.*"));
+        doReturn(CompletableFuture.completedFuture(Optional.empty())).when(zkDataCache).getAsync(matches(".*nonexistent.*"));
 
         AuthorizationManager authorizationManager = spy(new AuthorizationManager(svcConfig, configCacheService));
         doReturn(authorizationManager).when(brokerService).getAuthorizationManager();
@@ -440,7 +440,7 @@ public class ServerCnxTest {
         doReturn(authorizationManager).when(brokerService).getAuthorizationManager();
         doReturn(true).when(brokerService).isAuthorizationEnabled();
         doReturn(false).when(authorizationManager).isSuperUser(Mockito.anyString());
-        doReturn(true).when(authorizationManager).checkPermission(any(DestinationName.class), Mockito.anyString(),
+        doReturn(CompletableFuture.completedFuture(true)).when(authorizationManager).checkPermission(any(DestinationName.class), Mockito.anyString(),
                 any(AuthAction.class));
 
         resetChannel();
@@ -493,7 +493,7 @@ public class ServerCnxTest {
 
     public void testProducerCommandWithAuthorizationNegative() throws Exception {
         AuthorizationManager authorizationManager = mock(AuthorizationManager.class);
-        doReturn(false).when(authorizationManager).canProduce(Mockito.any(), Mockito.any());
+        doReturn(CompletableFuture.completedFuture(false)).when(authorizationManager).canProduceAsync(Mockito.any(), Mockito.any());
         doReturn(authorizationManager).when(brokerService).getAuthorizationManager();
         doReturn(true).when(brokerService).isAuthenticationEnabled();
         doReturn(true).when(brokerService).isAuthorizationEnabled();
@@ -1022,7 +1022,7 @@ public class ServerCnxTest {
     @Test(timeOut = 30000)
     public void testSubscribeCommandWithAuthorizationPositive() throws Exception {
         AuthorizationManager authorizationManager = mock(AuthorizationManager.class);
-        doReturn(true).when(authorizationManager).canConsume(Mockito.any(), Mockito.any());
+        doReturn(CompletableFuture.completedFuture(true)).when(authorizationManager).canConsumeAsync(Mockito.any(), Mockito.any());
         doReturn(authorizationManager).when(brokerService).getAuthorizationManager();
         doReturn(true).when(brokerService).isAuthenticationEnabled();
         doReturn(true).when(brokerService).isAuthorizationEnabled();
@@ -1042,7 +1042,7 @@ public class ServerCnxTest {
     @Test(timeOut = 30000)
     public void testSubscribeCommandWithAuthorizationNegative() throws Exception {
         AuthorizationManager authorizationManager = mock(AuthorizationManager.class);
-        doReturn(false).when(authorizationManager).canConsume(Mockito.any(), Mockito.any());
+        doReturn(CompletableFuture.completedFuture(false)).when(authorizationManager).canConsumeAsync(Mockito.any(), Mockito.any());
         doReturn(authorizationManager).when(brokerService).getAuthorizationManager();
         doReturn(true).when(brokerService).isAuthenticationEnabled();
         doReturn(true).when(brokerService).isAuthorizationEnabled();

pulsar-websocket/src/main/java/com/yahoo/pulsar/websocket/AbstractWebSocketHandler.java

@@ -22,6 +22,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.TreeMap;
+import java.util.concurrent.CompletableFuture;
 
 import javax.naming.AuthenticationException;
 import javax.servlet.http.HttpServletRequest;
@@ -73,11 +74,16 @@ public abstract class AbstractWebSocketHandler extends WebSocketAdapter implemen
             }
         }
 
-        if (service.isAuthorizationEnabled() && !isAuthorized(authRole)) {
-            log.warn("[{}] WebSocket Client [{}] is not authorized on topic {}", session.getRemoteAddress(), authRole,
-                    topic);
-            close(WebSocketError.NotAuthorizedError);
-            return;
+        if (service.isAuthorizationEnabled()) {
+            final String role = authRole;
+            isAuthorized(authRole).thenApply(isAuthorized -> {
+                if(!isAuthorized) {
+                    log.warn("[{}] WebSocket Client [{}] is not authorized on topic {}", session.getRemoteAddress(), role,
+                            topic);
+                    close(WebSocketError.NotAuthorizedError);
+                }
+                return null;
+            });
         }
     }
 
@@ -120,7 +126,7 @@ public abstract class AbstractWebSocketHandler extends WebSocketAdapter implemen
         return null;
     }
 
-    protected abstract boolean isAuthorized(String authRole);
+    protected abstract CompletableFuture<Boolean> isAuthorized(String authRole);
 
     private String extractTopicName(HttpServletRequest request) {
         String uri = request.getRequestURI();

pulsar-websocket/src/main/java/com/yahoo/pulsar/websocket/ConsumerHandler.java

@@ -24,6 +24,7 @@ import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
 import java.util.Base64;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -177,8 +178,8 @@ public class ConsumerHandler extends AbstractWebSocketHandler {
     }
 
     @Override
-    protected boolean isAuthorized(String authRole) {
-        return service.getAuthorizationManager().canConsume(DestinationName.get(topic), authRole);
+    protected CompletableFuture<Boolean> isAuthorized(String authRole) {
+        return service.getAuthorizationManager().canConsumeAsync(DestinationName.get(topic), authRole);
     }
 
     private static String extractSubscription(HttpServletRequest request) {

pulsar-websocket/src/main/java/com/yahoo/pulsar/websocket/ProducerHandler.java

@@ -17,6 +17,7 @@ package com.yahoo.pulsar.websocket;
 
 import java.io.IOException;
 import java.util.Base64;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
 import javax.servlet.http.HttpServletRequest;
@@ -119,8 +120,8 @@ public class ProducerHandler extends AbstractWebSocketHandler {
         });
     }
 
-    protected boolean isAuthorized(String authRole) {
-        return service.getAuthorizationManager().canProduce(DestinationName.get(topic), authRole);
+    protected CompletableFuture<Boolean> isAuthorized(String authRole) {
+        return service.getAuthorizationManager().canProduceAsync(DestinationName.get(topic), authRole);
     }
 
     private void sendAckResponse(ProducerAck response) {