Move impersonation authorization to SessionSupplier

Author: prithvip

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java

@@ -31,15 +31,13 @@
 import com.facebook.presto.server.SessionContext;
 import com.facebook.presto.server.SessionPropertyDefaults;
 import com.facebook.presto.server.SessionSupplier;
-import com.facebook.presto.server.security.SecurityConfig;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.QueryId;
 import com.facebook.presto.spi.analyzer.AnalyzerOptions;
 import com.facebook.presto.spi.analyzer.QueryPreparerProvider;
 import com.facebook.presto.spi.resourceGroups.SelectionContext;
 import com.facebook.presto.spi.resourceGroups.SelectionCriteria;
 import com.facebook.presto.spi.security.AccessControl;
-import com.facebook.presto.spi.security.AuthorizedIdentity;
 import com.facebook.presto.sql.analyzer.QueryPreparerProviderManager;
 import com.facebook.presto.transaction.TransactionManager;
 import com.google.common.util.concurrent.AbstractFuture;
@@ -57,8 +55,6 @@
 import java.util.concurrent.Executor;
 
 import static com.facebook.presto.SystemSessionProperties.getAnalyzerType;
-import static com.facebook.presto.security.AccessControlUtils.checkPermissions;
-import static com.facebook.presto.security.AccessControlUtils.getAuthorizedIdentity;
 import static com.facebook.presto.spi.StandardErrorCode.QUERY_TEXT_TOO_LARGE;
 import static com.facebook.presto.util.AnalyzerUtil.createAnalyzerOptions;
 import static com.google.common.base.Preconditions.checkArgument;
@@ -93,7 +89,6 @@ public class DispatchManager
 
     private final QueryManagerStats stats = new QueryManagerStats();
 
-    private final SecurityConfig securityConfig;
     private final QueryPreparerProviderManager queryPreparerProviderManager;
 
     /**
@@ -130,7 +125,6 @@ public DispatchManager(
             QueryManagerConfig queryManagerConfig,
             DispatchExecutor dispatchExecutor,
             ClusterStatusSender clusterStatusSender,
-            SecurityConfig securityConfig,
             Optional<ClusterQueryTrackerService> clusterQueryTrackerService)
     {
         this.queryIdGenerator = requireNonNull(queryIdGenerator, "queryIdGenerator is null");
@@ -152,8 +146,6 @@ public DispatchManager(
         this.clusterStatusSender = requireNonNull(clusterStatusSender, "clusterStatusSender is null");
 
         this.queryTracker = new QueryTracker<>(queryManagerConfig, dispatchExecutor.getScheduledExecutor(), clusterQueryTrackerService);
-
-        this.securityConfig = requireNonNull(securityConfig, "securityConfig is null");
     }
 
     /**
@@ -275,14 +267,8 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
                 throw new PrestoException(QUERY_TEXT_TOO_LARGE, format("Query text length (%s) exceeds the maximum length (%s)", queryLength, maxQueryLength));
             }
 
-            // check permissions if needed
-            checkPermissions(accessControl, securityConfig, queryId, sessionContext);
-
-            // get authorized identity if possible
-            Optional<AuthorizedIdentity> authorizedIdentity = getAuthorizedIdentity(accessControl, securityConfig, queryId, sessionContext);
-
             // decode session
-            session = sessionSupplier.createSession(queryId, sessionContext, warningCollectorFactory, authorizedIdentity);
+            session = sessionSupplier.createSession(queryId, sessionContext, warningCollectorFactory);
 
             // prepare query
             AnalyzerOptions analyzerOptions = createAnalyzerOptions(session, session.getWarningCollector());

presto-main/src/main/java/com/facebook/presto/server/NoOpSessionSupplier.java

@@ -16,9 +16,6 @@
 import com.facebook.presto.Session;
 import com.facebook.presto.execution.warnings.WarningCollectorFactory;
 import com.facebook.presto.spi.QueryId;
-import com.facebook.presto.spi.security.AuthorizedIdentity;
-
-import java.util.Optional;
 
 /**
  * Used on workers.
@@ -27,7 +24,7 @@ public class NoOpSessionSupplier
         implements SessionSupplier
 {
     @Override
-    public Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory, Optional<AuthorizedIdentity> authorizedIdentity)
+    public Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
     {
         throw new UnsupportedOperationException();
     }

presto-main/src/main/java/com/facebook/presto/server/QuerySessionSupplier.java

@@ -19,6 +19,7 @@
 import com.facebook.presto.common.type.TimeZoneKey;
 import com.facebook.presto.execution.warnings.WarningCollectorFactory;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.server.security.SecurityConfig;
 import com.facebook.presto.spi.QueryId;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.function.SqlFunctionId;
@@ -39,6 +40,8 @@
 import static com.facebook.presto.Session.SessionBuilder;
 import static com.facebook.presto.SystemSessionProperties.WARNING_HANDLING;
 import static com.facebook.presto.common.type.TimeZoneKey.getTimeZoneKey;
+import static com.facebook.presto.security.AccessControlUtils.checkPermissions;
+import static com.facebook.presto.security.AccessControlUtils.getAuthorizedIdentity;
 import static java.util.Map.Entry;
 import static java.util.Objects.requireNonNull;
 
@@ -52,44 +55,30 @@ public class QuerySessionSupplier
     private final AccessControl accessControl;
     private final SessionPropertyManager sessionPropertyManager;
     private final Optional<TimeZoneKey> forcedSessionTimeZone;
+    private final SecurityConfig securityConfig;
 
     @Inject
     public QuerySessionSupplier(
             TransactionManager transactionManager,
             AccessControl accessControl,
             SessionPropertyManager sessionPropertyManager,
-            SqlEnvironmentConfig config)
+            SqlEnvironmentConfig sqlEnvironmentConfig,
+            SecurityConfig securityConfig)
     {
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
         this.sessionPropertyManager = requireNonNull(sessionPropertyManager, "sessionPropertyManager is null");
-        requireNonNull(config, "config is null");
-        this.forcedSessionTimeZone = requireNonNull(config.getForcedSessionTimeZone(), "forcedSessionTimeZone is null");
+        requireNonNull(sqlEnvironmentConfig, "sqlEnvironmentConfig is null");
+        this.forcedSessionTimeZone = requireNonNull(sqlEnvironmentConfig.getForcedSessionTimeZone(), "forcedSessionTimeZone is null");
+        this.securityConfig = requireNonNull(securityConfig, "securityConfig is null");
     }
 
     @Override
-    public Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory, Optional<AuthorizedIdentity> authorizedIdentity)
+    public Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
     {
-        Identity identity = context.getIdentity();
-        if (authorizedIdentity.isPresent()) {
-            identity = new Identity(
-                    identity.getUser(),
-                    identity.getPrincipal(),
-                    identity.getRoles(),
-                    identity.getExtraCredentials(),
-                    identity.getExtraAuthenticators(),
-                    Optional.of(authorizedIdentity.get().getUserName()),
-                    authorizedIdentity.get().getReasonForSelect());
-            log.info(String.format(
-                    "For query %s, given user is %s, authorized user is %s",
-                    queryId.getId(),
-                    identity.getUser(),
-                    authorizedIdentity.get().getUserName()));
-        }
-
         SessionBuilder sessionBuilder = Session.builder(sessionPropertyManager)
                 .setQueryId(queryId)
-                .setIdentity(identity)
+                .setIdentity(authenticateIdentity(queryId, context))
                 .setSource(context.getSource())
                 .setCatalog(context.getCatalog())
                 .setSchema(context.getSchema())
@@ -145,4 +134,19 @@ else if (context.getTimeZoneId() != null) {
         }
         return session;
     }
+
+    private Identity authenticateIdentity(QueryId queryId, SessionContext context)
+    {
+        checkPermissions(accessControl, securityConfig, queryId, context);
+        Optional<AuthorizedIdentity> authorizedIdentity = getAuthorizedIdentity(accessControl, securityConfig, queryId, context);
+
+        return authorizedIdentity.map(identity -> new Identity(
+                context.getIdentity().getUser(),
+                context.getIdentity().getPrincipal(),
+                context.getIdentity().getRoles(),
+                context.getIdentity().getExtraCredentials(),
+                context.getIdentity().getExtraAuthenticators(),
+                Optional.of(identity.getUserName()),
+                identity.getReasonForSelect())).orElseGet(context::getIdentity);
+    }
 }

presto-main/src/main/java/com/facebook/presto/server/SessionSupplier.java

@@ -16,11 +16,8 @@
 import com.facebook.presto.Session;
 import com.facebook.presto.execution.warnings.WarningCollectorFactory;
 import com.facebook.presto.spi.QueryId;
-import com.facebook.presto.spi.security.AuthorizedIdentity;
-
-import java.util.Optional;
 
 public interface SessionSupplier
 {
-    Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory, Optional<AuthorizedIdentity> authorizedIdentity);
+    Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory);
 }

presto-main/src/test/java/com/facebook/presto/server/TestQuerySessionSupplier.java

@@ -18,6 +18,7 @@
 import com.facebook.presto.common.type.TimeZoneNotSupportedException;
 import com.facebook.presto.execution.warnings.WarningCollectorFactory;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.server.security.SecurityConfig;
 import com.facebook.presto.spi.QueryId;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.function.SqlFunctionId;
@@ -33,7 +34,6 @@
 import javax.servlet.http.HttpServletRequest;
 
 import java.util.Locale;
-import java.util.Optional;
 
 import static com.facebook.airlift.json.JsonCodec.jsonCodec;
 import static com.facebook.presto.SystemSessionProperties.HASH_PARTITION_COUNT;
@@ -91,7 +91,8 @@ public void testCreateSession()
                 createTestTransactionManager(),
                 new AllowAllAccessControl(),
                 new SessionPropertyManager(),
-                new SqlEnvironmentConfig());
+                new SqlEnvironmentConfig(),
+                new SecurityConfig());
         WarningCollectorFactory warningCollectorFactory = new WarningCollectorFactory()
         {
             @Override
@@ -100,7 +101,7 @@ public WarningCollector create(WarningHandlingLevel warningHandlingLevel)
                 return WarningCollector.NOOP;
             }
         };
-        Session session = sessionSupplier.createSession(new QueryId("test_query_id"), context, warningCollectorFactory, Optional.empty());
+        Session session = sessionSupplier.createSession(new QueryId("test_query_id"), context, warningCollectorFactory);
 
         assertEquals(session.getQueryId(), new QueryId("test_query_id"));
         assertEquals(session.getUser(), "testUser");
@@ -162,7 +163,8 @@ public void testInvalidTimeZone()
                 createTestTransactionManager(),
                 new AllowAllAccessControl(),
                 new SessionPropertyManager(),
-                new SqlEnvironmentConfig());
+                new SqlEnvironmentConfig(),
+                new SecurityConfig());
         WarningCollectorFactory warningCollectorFactory = new WarningCollectorFactory()
         {
             @Override
@@ -171,6 +173,6 @@ public WarningCollector create(WarningHandlingLevel warningHandlingLevel)
                 return WarningCollector.NOOP;
             }
         };
-        sessionSupplier.createSession(new QueryId("test_query_id"), context, warningCollectorFactory, Optional.empty());
+        sessionSupplier.createSession(new QueryId("test_query_id"), context, warningCollectorFactory);
     }
 }

presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java

@@ -91,7 +91,6 @@
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
 import com.facebook.presto.spi.security.AccessControl;
-import com.facebook.presto.spi.security.AuthorizedIdentity;
 import com.facebook.presto.spi.storage.TempStorage;
 import com.facebook.presto.sql.analyzer.BuiltInQueryPreparer;
 import com.facebook.presto.sql.analyzer.BuiltInQueryPreparer.BuiltInPreparedQuery;
@@ -139,8 +138,6 @@
 import static com.facebook.presto.execution.QueryState.FAILED;
 import static com.facebook.presto.execution.QueryState.PLANNING;
 import static com.facebook.presto.execution.StageInfo.getAllStages;
-import static com.facebook.presto.security.AccessControlUtils.checkPermissions;
-import static com.facebook.presto.security.AccessControlUtils.getAuthorizedIdentity;
 import static com.facebook.presto.server.protocol.QueryResourceUtil.toStatementStats;
 import static com.facebook.presto.spark.PrestoSparkSessionProperties.isAdaptiveQueryExecutionEnabled;
 import static com.facebook.presto.spark.SparkErrorCode.MALFORMED_QUERY_FILE;
@@ -612,13 +609,7 @@ public IPrestoSparkQueryExecution create(
                 credentialsProviders,
                 authenticatorProviders);
 
-        // check permissions if needed
-        checkPermissions(accessControl, securityConfig, queryId, sessionContext);
-
-        // get authorized identity if possible
-        Optional<AuthorizedIdentity> authorizedIdentity = getAuthorizedIdentity(accessControl, securityConfig, queryId, sessionContext);
-
-        Session session = sessionSupplier.createSession(queryId, sessionContext, warningCollectorFactory, authorizedIdentity);
+        Session session = sessionSupplier.createSession(queryId, sessionContext, warningCollectorFactory);
         session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, Optional.empty(), Optional.empty());
 
         if (!executionStrategies.isEmpty()) {