Support for groovy static analysis for groovy scripts

pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java

@@ -81,6 +81,7 @@
 import org.apache.pinot.core.util.ListenerConfigUtil;
 import org.apache.pinot.query.mailbox.MailboxService;
 import org.apache.pinot.query.service.dispatch.QueryDispatcher;
+import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
 import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
 import org.apache.pinot.spi.cursors.ResponseStoreService;
 import org.apache.pinot.spi.env.PinotConfiguration;
@@ -478,6 +479,10 @@ public void start()
     _participantHelixManager.addPreConnectCallback(
         () -> _brokerMetrics.addMeteredGlobalValue(BrokerMeter.HELIX_ZOOKEEPER_RECONNECTS, 1L));
 
+    // Initializing Groovy execution security
+    GroovyFunctionEvaluator.configureGroovySecurity(
+        _brokerConf.getProperty(CommonConstants.GROOVY_STATIC_ANALYZER_CONFIG));
+
     // Register the service status handler
     registerServiceStatusHandler();
 

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java

@@ -90,6 +90,7 @@
 import org.apache.pinot.core.transport.ServerInstance;
 import org.apache.pinot.core.util.GapfillUtils;
 import org.apache.pinot.query.parser.utils.ParserUtils;
+import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
 import org.apache.pinot.spi.auth.AuthorizationResult;
 import org.apache.pinot.spi.config.table.FieldConfig;
 import org.apache.pinot.spi.config.table.QueryConfig;
@@ -515,9 +516,7 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn
     }
 
     HandlerContext handlerContext = getHandlerContext(offlineTableConfig, realtimeTableConfig);
-    if (handlerContext._disableGroovy) {
-      rejectGroovyQuery(serverPinotQuery);
-    }
+    rejectGroovyQuery(serverPinotQuery, handlerContext._disableGroovy);
     if (handlerContext._useApproximateFunction) {
       handleApproximateFunctionOverride(serverPinotQuery);
     }
@@ -1429,45 +1428,59 @@ private static class HandlerContext {
    * Verifies that no groovy is present in the PinotQuery when disabled.
    */
   @VisibleForTesting
-  static void rejectGroovyQuery(PinotQuery pinotQuery) {
+  static void rejectGroovyQuery(PinotQuery pinotQuery, boolean disableGroovy) {
     List<Expression> selectList = pinotQuery.getSelectList();
     for (Expression expression : selectList) {
-      rejectGroovyQuery(expression);
+      rejectGroovyQuery(expression, disableGroovy);
     }
     List<Expression> orderByList = pinotQuery.getOrderByList();
     if (orderByList != null) {
       for (Expression expression : orderByList) {
         // NOTE: Order-by is always a Function with the ordering of the Expression
-        rejectGroovyQuery(expression.getFunctionCall().getOperands().get(0));
+        rejectGroovyQuery(expression.getFunctionCall().getOperands().get(0), disableGroovy);
       }
     }
     Expression havingExpression = pinotQuery.getHavingExpression();
     if (havingExpression != null) {
-      rejectGroovyQuery(havingExpression);
+      rejectGroovyQuery(havingExpression, disableGroovy);
     }
     Expression filterExpression = pinotQuery.getFilterExpression();
     if (filterExpression != null) {
-      rejectGroovyQuery(filterExpression);
+      rejectGroovyQuery(filterExpression, disableGroovy);
     }
     List<Expression> groupByList = pinotQuery.getGroupByList();
     if (groupByList != null) {
       for (Expression expression : groupByList) {
-        rejectGroovyQuery(expression);
+        rejectGroovyQuery(expression, disableGroovy);
       }
     }
   }
 
-  private static void rejectGroovyQuery(Expression expression) {
+  private static void rejectGroovyQuery(Expression expression, boolean disableGroovy) {
     Function function = expression.getFunctionCall();
     if (function == null) {
       return;
     }
     if (function.getOperator().equals("groovy")) {
-      throw new BadQueryRequestException("Groovy transform functions are disabled for queries");
+      if (disableGroovy) {
+        throw new BadQueryRequestException("Groovy transform functions are disabled for queries");
+      } else {
+        groovySecureAnalysis(function);
+      }
     }
     for (Expression operandExpression : function.getOperands()) {
-      rejectGroovyQuery(operandExpression);
+      rejectGroovyQuery(operandExpression, disableGroovy);
+    }
+  }
+
+  private static void groovySecureAnalysis(Function function) {
+    List<Expression> operands = function.getOperands();
+    if (operands == null || operands.size() < 2) {
+      throw new BadQueryRequestException("Groovy transform function must have at least 2 argument");
     }
+    // second argument in the groovy function is groovy script
+    String script = operands.get(1).getLiteral().getStringValue();
+    GroovyFunctionEvaluator.parseGroovyScript(String.format("groovy({%s})", script));
   }
 
   /**

pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/QueryValidationTest.java

@@ -19,13 +19,19 @@
 
 package org.apache.pinot.broker.requesthandler;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.common.collect.ImmutableMap;
 import java.util.Map;
 import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
+import org.apache.pinot.segment.local.function.GroovyStaticAnalyzerConfig;
 import org.apache.pinot.sql.parsers.CalciteSqlParser;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+
 
 public class QueryValidationTest {
 
@@ -113,6 +119,46 @@ public void testRejectGroovyQuery() {
     testRejectGroovyQuery("SELECT foo FROM bar", false);
   }
 
+  @Test
+  public void testGroovyScripts()
+      throws JsonProcessingException {
+    // setup secure groovy config
+    GroovyFunctionEvaluator.setGroovyStaticAnalyzerConfig(GroovyStaticAnalyzerConfig.createDefault());
+
+    String inValidGroovyQuery = "SELECT groovy('{\"returnType\":\"INT\",\"isSingleValue\":true}') FROM foo";
+    runUnsupportedGroovy(inValidGroovyQuery, "Groovy transform function must have at least 2 argument");
+
+    String groovyInvalidMethodInvokeQuery =
+        "SELECT groovy('{\"returnType\":\"STRING\",\"isSingleValue\":true}', 'return [\"bash\", \"-c\", \"echo Hello,"
+            + " World!\"].execute().text') FROM foo";
+    runUnsupportedGroovy(groovyInvalidMethodInvokeQuery, "Expression [MethodCallExpression] is not allowed");
+
+    String groovyInvalidImportsQuery =
+        "SELECT groovy( '{\"returnType\":\"INT\",\"isSingleValue\":true}', 'def args = [\"QuickStart\", \"-type\", "
+            + "\"REALTIME\"] as String[]; org.apache.pinot.tools.admin.PinotAdministrator.main(args); 2') FROM foo";
+    runUnsupportedGroovy(groovyInvalidImportsQuery, "Indirect import checks prevents usage of expression");
+
+    String groovyInOrderByClause =
+        "SELECT colA, colB FROM foo ORDER BY groovy('{\"returnType\":\"STRING\",\"isSingleValue\":true}', 'return "
+            + "[\"bash\", \"-c\", \"echo Hello, World!\"].execute().text') DESC";
+    runUnsupportedGroovy(groovyInOrderByClause, "Expression [MethodCallExpression] is not allowed");
+
+    String groovyInHavingClause =
+        "SELECT colA, SUM(colB) AS totalB, groovy('{\"returnType\":\"DOUBLE\",\"isSingleValue\":true}', 'arg0 / "
+            + "arg1', SUM(colB), COUNT(*)) AS avgB FROM foo GROUP BY colA HAVING groovy('{\"returnType\":\"BOOLEAN\","
+            + "\"isSingleValue\":true}', 'System.metaClass.methods.each { method -> if (method.name.md5() == "
+            + "\"f24f62eeb789199b9b2e467df3b1876b\") {method.invoke(System, 10)} }', SUM(colB))";
+    runUnsupportedGroovy(groovyInHavingClause, "Indirect import checks prevents usage of expression");
+
+    String groovyInWhereClause =
+        "SELECT colA, colB FROM foo WHERE groovy('{\"returnType\":\"BOOLEAN\",\"isSingleValue\":true}', 'System.exit"
+            + "(10)', colA)";
+    runUnsupportedGroovy(groovyInWhereClause, "Indirect import checks prevents usage of expression");
+
+    // Reset groovy config for rest of the testing
+    GroovyFunctionEvaluator.setGroovyStaticAnalyzerConfig(null);
+  }
+
   @Test
   public void testReplicaGroupToQueryInvalidQuery() {
     PinotQuery pinotQuery =
@@ -125,20 +171,30 @@ private void testRejectGroovyQuery(String query, boolean queryContainsGroovy) {
     PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
 
     try {
-      BaseSingleStageBrokerRequestHandler.rejectGroovyQuery(pinotQuery);
+      BaseSingleStageBrokerRequestHandler.rejectGroovyQuery(pinotQuery, queryContainsGroovy);
       if (queryContainsGroovy) {
-        Assert.fail("Query should have failed since groovy was found in query: " + pinotQuery);
+        fail("Query should have failed since groovy was found in query: " + pinotQuery);
       }
     } catch (Exception e) {
       Assert.assertEquals(e.getMessage(), "Groovy transform functions are disabled for queries");
     }
   }
 
+  private static void runUnsupportedGroovy(String query, String errorMsg) {
+    try {
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+      BaseSingleStageBrokerRequestHandler.rejectGroovyQuery(pinotQuery, false);
+      fail("Query should have failed since malicious groovy was found in query");
+    } catch (Exception e) {
+      assertTrue(e.getMessage().contains(errorMsg));
+    }
+  }
+
   private void testUnsupportedQuery(String query, String errorMessage) {
     try {
       PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
       BaseSingleStageBrokerRequestHandler.validateRequest(pinotQuery, 1000);
-      Assert.fail("Query should have failed");
+      fail("Query should have failed");
     } catch (Exception e) {
       Assert.assertEquals(e.getMessage(), errorMessage);
     }
@@ -149,7 +205,7 @@ private void testNonExistingColumns(String rawTableName, boolean isCaseInsensiti
     try {
       PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
       BaseSingleStageBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap);
-      Assert.fail("Query should have failed");
+      fail("Query should have failed");
     } catch (Exception e) {
       Assert.assertEquals(errorMessage, e.getMessage());
     }
@@ -161,7 +217,7 @@ private void testExistingColumns(String rawTableName, boolean isCaseInsensitive,
       PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
       BaseSingleStageBrokerRequestHandler.updateColumnNames(rawTableName, pinotQuery, isCaseInsensitive, columnNameMap);
     } catch (Exception e) {
-      Assert.fail("Query should have succeeded");
+      fail("Query should have succeeded");
     }
   }
 }

pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java

@@ -70,7 +70,6 @@ public enum ControllerMeter implements AbstractMetrics.Meter {
   IDEAL_STATE_UPDATE_RETRY("IdealStateUpdateRetry", false),
   IDEAL_STATE_UPDATE_SUCCESS("IdealStateUpdateSuccess", false);
 
-
   private final String _brokerMeterName;
   private final String _unit;
   private final boolean _global;

pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java

@@ -125,6 +125,7 @@
 import org.apache.pinot.core.segment.processing.lifecycle.PinotSegmentLifecycleEventListenerManager;
 import org.apache.pinot.core.transport.ListenerConfig;
 import org.apache.pinot.core.util.ListenerConfigUtil;
+import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
 import org.apache.pinot.segment.local.utils.TableConfigUtils;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.crypt.PinotCrypterFactory;
@@ -387,7 +388,8 @@ public ControllerConf getConfig() {
   }
 
   @Override
-  public void start() {
+  public void start()
+      throws Exception {
     LOGGER.info("Starting Pinot controller in mode: {}. (Version: {})", _controllerMode.name(), PinotVersion.VERSION);
     LOGGER.info("Controller configs: {}", new PinotAppConfigs(getConfig()).toJSONString());
     long startTimeMs = System.currentTimeMillis();
@@ -412,6 +414,10 @@ public void start() {
         break;
     }
 
+    // Initializing Groovy execution security
+    GroovyFunctionEvaluator.configureGroovySecurity(
+        _config.getProperty(CommonConstants.GROOVY_STATIC_ANALYZER_CONFIG));
+
     ServiceStatus.setServiceStatusCallback(_helixParticipantInstanceId,
         new ServiceStatus.MultipleCallbackServiceStatusCallback(_serviceStatusCallbackList));
     _controllerMetrics.addTimedValue(ControllerTimer.STARTUP_SUCCESS_DURATION_MS,

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotClusterConfigs.java

@@ -55,6 +55,9 @@
 import org.apache.pinot.core.auth.Actions;
 import org.apache.pinot.core.auth.Authorize;
 import org.apache.pinot.core.auth.TargetType;
+import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
+import org.apache.pinot.segment.local.function.GroovyStaticAnalyzerConfig;
+import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.JsonUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -168,4 +171,73 @@ public SuccessResponse deleteClusterConfig(
       throw new ControllerApplicationException(LOGGER, errStr, Response.Status.INTERNAL_SERVER_ERROR, e);
     }
   }
+
+  @GET
+  @Path("/cluster/configs/groovy/staticAnalyzerConfig")
+  @Authorize(targetType = TargetType.CLUSTER, action = Actions.Cluster.GET_GROOVY_STATIC_ANALYZER_CONFIG)
+  @Produces(MediaType.APPLICATION_JSON)
+  @ApiOperation(value = "Get the configuration for Groovy Static analysis",
+      notes = "Get the configuration for Groovy static analysis")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Success"),
+      @ApiResponse(code = 500, message = "Internal server error")
+  })
+  public GroovyStaticAnalyzerConfig getGroovyStaticAnalysisConfig() throws Exception {
+    HelixAdmin helixAdmin = _pinotHelixResourceManager.getHelixAdmin();
+    HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER)
+        .forCluster(_pinotHelixResourceManager.getHelixClusterName()).build();
+    Map<String, String> configs = helixAdmin.getConfig(configScope,
+        List.of(CommonConstants.GROOVY_STATIC_ANALYZER_CONFIG));
+    String json = configs.get(CommonConstants.GROOVY_STATIC_ANALYZER_CONFIG);
+    if (json != null) {
+      return GroovyStaticAnalyzerConfig.fromJson(json);
+    } else {
+      return null;
+    }
+  }
+
+  @POST
+  @Path("/cluster/configs/groovy/staticAnalyzerConfig")
+  @Authorize(targetType = TargetType.CLUSTER, action = Actions.Cluster.UPDATE_GROOVY_STATIC_ANALYZER_CONFIG)
+  @Authenticate(AccessType.UPDATE)
+  @ApiOperation(value = "Update Groovy static analysis configuration")
+  @Produces(MediaType.APPLICATION_JSON)
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Success"),
+      @ApiResponse(code = 500, message = "Server error updating configuration")
+  })
+  public SuccessResponse setGroovyStaticAnalysisConfig(String body) throws Exception {
+    try {
+      HelixAdmin admin = _pinotHelixResourceManager.getHelixAdmin();
+      HelixConfigScope configScope =
+          new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+              _pinotHelixResourceManager.getHelixClusterName()).build();
+      Map<String, String> properties = new TreeMap<>();
+      GroovyStaticAnalyzerConfig groovyConfig = GroovyStaticAnalyzerConfig.fromJson(body);
+      properties.put(CommonConstants.GROOVY_STATIC_ANALYZER_CONFIG, groovyConfig.toJson());
+      admin.setConfig(configScope, properties);
+      GroovyFunctionEvaluator.setGroovyStaticAnalyzerConfig(groovyConfig);
+      return new SuccessResponse("Updated Groovy Static Analyzer config.");
+    } catch (IOException e) {
+      throw new ControllerApplicationException(LOGGER, "Error converting request to cluster config",
+          Response.Status.BAD_REQUEST, e);
+    } catch (Exception e) {
+      throw new ControllerApplicationException(LOGGER, "Failed to update Groovy Static Analyzer config",
+          Response.Status.INTERNAL_SERVER_ERROR, e);
+    }
+  }
+
+  @GET
+  @Path("/cluster/configs/groovy/staticAnalyzerConfig/default")
+  @Authorize(targetType = TargetType.CLUSTER, action = Actions.Cluster.GET_GROOVY_STATIC_ANALYZER_CONFIG)
+  @Produces(MediaType.APPLICATION_JSON)
+  @ApiOperation(value = "Get the default configuration for Groovy Static analysis",
+      notes = "Get the default configuration for Groovy static analysis")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Success"),
+      @ApiResponse(code = 500, message = "Internal server error")
+  })
+  public GroovyStaticAnalyzerConfig getDefaultGroovyStaticAnalysisConfig() {
+    return GroovyStaticAnalyzerConfig.createDefault();
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/auth/Actions.java

@@ -99,6 +99,8 @@ public static class Cluster {
     public static final String UPDATE_INSTANCE_PARTITIONS = "UpdateInstancePartitions";
     public static final String GET_RESPONSE_STORE = "GetResponseStore";
     public static final String DELETE_RESPONSE_STORE = "DeleteResponseStore";
+    public static final String GET_GROOVY_STATIC_ANALYZER_CONFIG = "GetGroovyStaticAnalyzerConfig";
+    public static final String UPDATE_GROOVY_STATIC_ANALYZER_CONFIG = "UpdateGroovyStaticAnalyzerConfig";
   }
 
   // Action names for table