Support for groovy static analysis for groovy scripts

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

@@ -68,7 +68,9 @@ public enum ControllerMeter implements AbstractMetrics.Meter {
   NUMBER_ADHOC_TASKS_SUBMITTED("adhocTasks", false),
   IDEAL_STATE_UPDATE_FAILURE("IdealStateUpdateFailure", false),
   IDEAL_STATE_UPDATE_RETRY("IdealStateUpdateRetry", false),
-  IDEAL_STATE_UPDATE_SUCCESS("IdealStateUpdateSuccess", false);
+  IDEAL_STATE_UPDATE_SUCCESS("IdealStateUpdateSuccess", false),
+  GROOVY_SECURITY_VIOLATIONS("exceptions", true);
+
 
 
   private final String _brokerMeterName;

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

@@ -38,7 +38,8 @@ public enum MinionMeter implements AbstractMetrics.Meter {
   SEGMENT_BYTES_UPLOADED("bytes", false),
   RECORDS_PROCESSED_COUNT("rows", false),
   RECORDS_PURGED_COUNT("rows", false),
-  COMPACTED_RECORDS_COUNT("rows", false);
+  COMPACTED_RECORDS_COUNT("rows", false),
+  GROOVY_SECURITY_VIOLATIONS("exceptions", true);
 
   private final String _meterName;
   private final String _unit;

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

@@ -27,6 +27,7 @@
 public enum ServerMeter implements AbstractMetrics.Meter {
   QUERIES("queries", true),
   UNCAUGHT_EXCEPTIONS("exceptions", true),
+  GROOVY_SECURITY_VIOLATIONS("exceptions", true),
   REQUEST_DESERIALIZATION_EXCEPTIONS("exceptions", true),
   RESPONSE_SERIALIZATION_EXCEPTIONS("exceptions", true),
   SCHEDULING_TIMEOUT_EXCEPTIONS("exceptions", true),

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

@@ -122,6 +122,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;
@@ -382,7 +383,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();
@@ -407,6 +409,12 @@ public void start() {
         break;
     }
 
+    // Initializing Groovy execution security
+    GroovyFunctionEvaluator.configureGroovySecurity(
+        _config.getProperty(CommonConstants.GROOVY_STATIC_ANALYZER_CONFIG));
+    GroovyFunctionEvaluator.setMetrics(_controllerMetrics, ControllerMeter.GROOVY_SECURITY_VIOLATIONS);
+
+
     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,75 @@ 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 == null ? null : groovyConfig.toJson());
+      admin.setConfig(configScope, properties);
+      GroovyFunctionEvaluator.setConfig(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

pinot-core/src/test/java/org/apache/pinot/core/data/function/GroovyFunctionEvaluatorTest.java

@@ -18,32 +18,123 @@
  */
 package org.apache.pinot.core.data.function;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
+import org.apache.pinot.segment.local.function.GroovyStaticAnalyzerConfig;
 import org.apache.pinot.spi.data.readers.GenericRow;
-import org.testng.Assert;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 import org.testng.collections.Lists;
 
+import static org.apache.pinot.segment.local.function.GroovyStaticAnalyzerConfig.getDefaultAllowedImports;
+import static org.apache.pinot.segment.local.function.GroovyStaticAnalyzerConfig.getDefaultAllowedReceivers;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.fail;
+
 
 /**
  * Tests Groovy functions for transforming schema columns
  */
 public class GroovyFunctionEvaluatorTest {
+  @Test
+  public void testLegalGroovyScripts()
+      throws JsonProcessingException {
+    // TODO: Add separate tests for these rules: receivers, imports, static imports, and method names.
+    List<String> scripts = List.of(
+        "Groovy({2})",
+        "Groovy({![\"pinot_minion_totalOutputSegmentSize_Value\"].contains(\"\");2})",
+        "Groovy({airtime == null ? (arrdelay == null ? 0 : arrdelay.value) : airtime.value; 2}, airtime, arrdelay)"
+    );
+
+    GroovyStaticAnalyzerConfig config = new GroovyStaticAnalyzerConfig(
+        getDefaultAllowedReceivers(),
+        getDefaultAllowedImports(),
+        getDefaultAllowedImports(),
+        List.of("invoke", "execute"),
+        false);
+    GroovyFunctionEvaluator.setConfig(config);
+
+    for (String script : scripts) {
+      GroovyFunctionEvaluator groovyFunctionEvaluator = new GroovyFunctionEvaluator(script);
+      GenericRow row = new GenericRow();
+      Object result = groovyFunctionEvaluator.evaluate(row);
+      assertEquals(2, result);
+    }
+  }
+
+  @Test
+  public void testIllegalGroovyScripts()
+      throws JsonProcessingException {
+    // TODO: Add separate tests for these rules: receivers, imports, static imports, and method names.
+    List<String> scripts = List.of(
+        "Groovy({\"ls\".execute()})",
+        "Groovy({[\"ls\"].execute()})",
+        "Groovy({System.exit(5)})",
+        "Groovy({System.metaClass.methods.each { method -> if (method.name.md5() == "
+            + "\"f24f62eeb789199b9b2e467df3b1876b\") {method.invoke(System, 10)} }})",
+        "Groovy({System.metaClass.methods.each { method -> if (method.name.reverse() == (\"ti\" + \"xe\")) "
+            + "{method.invoke(System, 10)} }})",
+        "groovy({def args = [\"QuickStart\", \"-type\", \"REALTIME\"] as String[]; "
+            + "org.apache.pinot.tools.admin.PinotAdministrator.main(args); 2})",
+        "Groovy({return [\"bash\", \"-c\", \"env\"].execute().text})"
+    );
+
+    GroovyStaticAnalyzerConfig config = new GroovyStaticAnalyzerConfig(
+        getDefaultAllowedReceivers(),
+        getDefaultAllowedImports(),
+        getDefaultAllowedImports(),
+        List.of("invoke", "execute"),
+        false);
+    GroovyFunctionEvaluator.setConfig(config);
+
+    for (String script : scripts) {
+      try {
+        new GroovyFunctionEvaluator(script);
+        fail("Groovy analyzer failed to catch malicious script");
+      } catch (Exception ignored) {
+      }
+    }
+  }
+
+  @Test
+  public void testUpdatingConfiguration()
+      throws JsonProcessingException {
+    // TODO: Figure out how to test this with the singleton initializer
+    // These tests would pass by default but the configuration will be updated so that they fail
+    List<String> scripts = List.of(
+        "Groovy({2})",
+        "Groovy({![\"pinot_minion_totalOutputSegmentSize_Value\"].contains(\"\");2})",
+        "Groovy({airtime == null ? (arrdelay == null ? 0 : arrdelay.value) : airtime.value; 2}, airtime, arrdelay)"
+    );
+
+    GroovyStaticAnalyzerConfig config =
+        new GroovyStaticAnalyzerConfig(List.of(), List.of(), List.of(), List.of(), false);
+    GroovyFunctionEvaluator.setConfig(config);
+
+    for (String script : scripts) {
+      try {
+        GroovyFunctionEvaluator groovyFunctionEvaluator = new GroovyFunctionEvaluator(script);
+        GenericRow row = new GenericRow();
+        groovyFunctionEvaluator.evaluate(row);
+        fail(String.format("Groovy analyzer failed to catch malicious script: %s", script));
+      } catch (Exception ignored) {
+      }
+    }
+  }
 
   @Test(dataProvider = "groovyFunctionEvaluationDataProvider")
   public void testGroovyFunctionEvaluation(String transformFunction, List<String> arguments, GenericRow genericRow,
       Object expectedResult) {
 
     GroovyFunctionEvaluator groovyExpressionEvaluator = new GroovyFunctionEvaluator(transformFunction);
-    Assert.assertEquals(groovyExpressionEvaluator.getArguments(), arguments);
+    assertEquals(groovyExpressionEvaluator.getArguments(), arguments);
 
     Object result = groovyExpressionEvaluator.evaluate(genericRow);
-    Assert.assertEquals(result, expectedResult);
+    assertEquals(result, expectedResult);
   }
 
   @DataProvider(name = "groovyFunctionEvaluationDataProvider")
@@ -108,20 +199,26 @@ public Object[][] groovyFunctionEvaluationDataProvider() {
     GenericRow genericRow9 = new GenericRow();
     genericRow9.putValue("ArrTime", 101);
     genericRow9.putValue("ArrTimeV2", null);
-    entries.add(new Object[]{"Groovy({ArrTimeV2 != null ? ArrTimeV2: ArrTime }, ArrTime, ArrTimeV2)",
-        Lists.newArrayList("ArrTime", "ArrTimeV2"), genericRow9, 101});
+    entries.add(new Object[]{
+        "Groovy({ArrTimeV2 != null ? ArrTimeV2: ArrTime }, ArrTime, ArrTimeV2)",
+        Lists.newArrayList("ArrTime", "ArrTimeV2"), genericRow9, 101
+    });
 
     GenericRow genericRow10 = new GenericRow();
     String jello = "Jello";
     genericRow10.putValue("jello", jello);
-    entries.add(new Object[]{"Groovy({jello != null ? jello.length() : \"Jello\" }, jello)",
-        Lists.newArrayList("jello"), genericRow10, 5});
+    entries.add(new Object[]{
+        "Groovy({jello != null ? jello.length() : \"Jello\" }, jello)",
+        Lists.newArrayList("jello"), genericRow10, 5
+    });
 
     //Invalid groovy script
     GenericRow genericRow11 = new GenericRow();
     genericRow11.putValue("nullValue", null);
-    entries.add(new Object[]{"Groovy({nullValue == null ? nullValue.length() : \"Jello\" }, nullValue)",
-        Lists.newArrayList("nullValue"), genericRow11, null});
+    entries.add(new Object[]{
+        "Groovy({nullValue == null ? nullValue.length() : \"Jello\" }, nullValue)",
+        Lists.newArrayList("nullValue"), genericRow11, null
+    });
     return entries.toArray(new Object[entries.size()][]);
   }
 }