From 08514a859d2cb9ebd8ca2c4c280d2f78e86a12e9 Mon Sep 17 00:00:00 2001 From: Soumitra Kumar Date: Wed, 18 Oct 2023 11:34:52 -0700 Subject: [PATCH] Adding RBAC authorization checks for multi-stage query engine --- .../MultiStageBrokerRequestHandler.java | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 5c4e86a7d4f9..772f021fce14 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -30,6 +30,7 @@ import javax.ws.rs.core.HttpHeaders; import javax.ws.rs.core.Response; import org.apache.calcite.jdbc.CalciteSchemaBuilder; +import org.apache.pinot.broker.api.AccessControl; import org.apache.pinot.broker.api.RequesterIdentity; import org.apache.pinot.broker.broker.AccessControlFactory; import org.apache.pinot.broker.querylog.QueryLogger; @@ -51,6 +52,8 @@ import org.apache.pinot.common.utils.ExceptionUtils; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.common.utils.request.RequestUtils; +import org.apache.pinot.core.auth.Actions; +import org.apache.pinot.core.auth.TargetType; import org.apache.pinot.core.query.reduce.ExecutionStatsAggregator; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.query.QueryEnvironment; @@ -128,7 +131,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S queryPlanResult = _queryEnvironment.explainQuery(query, sqlNodeAndOptions, requestId); String plan = queryPlanResult.getExplainPlan(); Set tableNames = queryPlanResult.getTableNames(); - if (!hasTableAccess(requesterIdentity, tableNames, requestContext)) { + if (!hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders)) { throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN); } @@ -164,7 +167,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S updatePhaseTimingForTables(tableNames, BrokerQueryPhase.REQUEST_COMPILATION, compilationTimeNs); // Validate table access. - if (!hasTableAccess(requesterIdentity, tableNames, requestContext)) { + if (!hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders)) { throw new WebApplicationException("Permission denied", Response.Status.FORBIDDEN); } updatePhaseTimingForTables(tableNames, BrokerQueryPhase.AUTHORIZATION, System.nanoTime() - compilationEndTimeNs); @@ -251,8 +254,10 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S * Validates whether the requester has access to all the tables. */ private boolean hasTableAccess(RequesterIdentity requesterIdentity, Set tableNames, - RequestContext requestContext) { - boolean hasAccess = _accessControlFactory.create().hasAccess(requesterIdentity, tableNames); + RequestContext requestContext, HttpHeaders httpHeaders) { + AccessControl accessControl = _accessControlFactory.create(); + boolean hasAccess = accessControl.hasAccess(requesterIdentity, tableNames) && tableNames.stream() + .allMatch(table -> accessControl.hasAccess(httpHeaders, TargetType.TABLE, table, Actions.Table.QUERY)); if (!hasAccess) { _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1); LOGGER.warn("Access denied for requestId {}", requestContext.getRequestId());