Clean up whitelist loading for core contexts in painless plugin (#76392)

This change creates a standard for loading additional allow listed elements for the core contexts 
within the ScriptPlugin. For each core context we look for org.elasticsearch.script.<context-name>.txt 
and if it exists we load it automatically for that context.

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessPlugin.java

@@ -38,19 +38,10 @@
 import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
-import org.elasticsearch.script.AggregationScript;
-import org.elasticsearch.script.FieldScript;
-import org.elasticsearch.script.FilterScript;
-import org.elasticsearch.script.IngestScript;
-import org.elasticsearch.script.NumberSortScript;
-import org.elasticsearch.script.ScoreScript;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
 import org.elasticsearch.script.ScriptModule;
 import org.elasticsearch.script.ScriptService;
-import org.elasticsearch.script.ScriptedMetricAggContexts;
-import org.elasticsearch.script.StringSortScript;
-import org.elasticsearch.search.aggregations.pipeline.MovingFunctionScript;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.watcher.ResourceWatcherService;
 
@@ -96,88 +87,28 @@ public final class PainlessPlugin extends Plugin implements ScriptPlugin, Extens
      * under Painless' resources
      */
     static {
-        Map<ScriptContext<?>, List<Whitelist>> map = new HashMap<>();
-
-        // Moving Function Pipeline Agg
-        List<Whitelist> movFn = new ArrayList<>();
-        Whitelist movFnWhitelist = WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.aggs.movfn.txt");
-        movFn.add(movFnWhitelist);
-        map.put(MovingFunctionScript.CONTEXT, movFn);
-
-        // Functions used for scoring docs
-        List<Whitelist> scoreFn = new ArrayList<>();
-        Whitelist scoreFnWhitelist = WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.score.txt");
-        Whitelist scoreFieldWhitelist =
-            WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.score.txt");
-        scoreFn.add(scoreFnWhitelist);
-        scoreFn.add(scoreFieldWhitelist);
-        map.put(ScoreScript.CONTEXT, scoreFn);
-
-        // Functions available to ingest pipelines
-        List<Whitelist> ingest = new ArrayList<>();
-        Whitelist ingestWhitelist = WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.ingest.txt");
-        ingest.add(ingestWhitelist);
-        map.put(IngestScript.CONTEXT, ingest);
-
-        // Functions available to runtime fields
-
-        for (ScriptContext<?> scriptContext : ScriptModule.RUNTIME_FIELDS_CONTEXTS) {
-            map.put(scriptContext, getRuntimeFieldWhitelist(scriptContext.name));
-        }
+        whitelists = new HashMap<>();
+
+        for (ScriptContext<?> context : ScriptModule.CORE_CONTEXTS.values()) {
+            List<Whitelist> contextWhitelists = new ArrayList<>();
+            if (PainlessPlugin.class.getResourceAsStream("org.elasticsearch.script." + context.name.replace('-', '_') + ".txt") != null) {
+                contextWhitelists.add(
+                        WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class,
+                                "org.elasticsearch.script." + context.name.replace('-', '_') + ".txt")
+                );
+            }
 
-        List<Whitelist> numSort = new ArrayList<>();
-        Whitelist numSortField =
-            WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.numbersort.txt");
-        numSort.add(numSortField);
-        map.put(NumberSortScript.CONTEXT, numSort);
-
-        List<Whitelist> strSort = new ArrayList<>();
-        Whitelist strSortField =
-            WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.stringsort.txt");
-        strSort.add(strSortField);
-        map.put(StringSortScript.CONTEXT, strSort);
-
-        List<Whitelist> filter = new ArrayList<>();
-        Whitelist filterWhitelist =
-            WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.filter.txt");
-        filter.add(filterWhitelist);
-        map.put(FilterScript.CONTEXT, filter);
-
-        List<Whitelist> aggregation = new ArrayList<>();
-        Whitelist aggregationWhitelist =
-                WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.aggregation.txt");
-        aggregation.add(aggregationWhitelist);
-        map.put(AggregationScript.CONTEXT, aggregation);
-
-        List<Whitelist> aggmap = new ArrayList<>();
-        Whitelist aggmapWhitelist =
-                WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.aggmap.txt");
-        aggmap.add(aggmapWhitelist);
-        map.put(ScriptedMetricAggContexts.MapScript.CONTEXT, aggmap);
-
-        List<Whitelist> field = new ArrayList<>();
-        Whitelist fieldWhitelist =
-                WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.field.txt");
-        field.add(fieldWhitelist);
-        map.put(FieldScript.CONTEXT, field);
-
-        // Execute context gets everything
-        List<Whitelist> test = new ArrayList<>();
-        test.add(movFnWhitelist);
-        test.add(scoreFnWhitelist);
-        test.add(ingestWhitelist);
-        test.add(WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.json.txt"));
-        map.put(PainlessExecuteAction.PainlessTestScript.CONTEXT, test);
-
-        whitelists = map;
-    }
+            whitelists.put(context, contextWhitelists);
+        }
 
-    private static List<Whitelist> getRuntimeFieldWhitelist(String contextName) {
-        List<Whitelist> scriptField = new ArrayList<>();
-        Whitelist whitelist = WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class,
-            "org.elasticsearch.script." + contextName + ".txt");
-        scriptField.add(whitelist);
-        return scriptField;
+        List<Whitelist> testWhitelists = new ArrayList<>();
+        for (ScriptContext<?> context : ScriptModule.CORE_CONTEXTS.values()) {
+            if (ScriptModule.RUNTIME_FIELDS_CONTEXTS.contains(context) == false) {
+                testWhitelists.addAll(whitelists.get(context));
+            }
+        }
+        testWhitelists.add(WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.json.txt"));
+        whitelists.put(PainlessExecuteAction.PainlessTestScript.CONTEXT, testWhitelists);
     }
 
     private final SetOnce<PainlessScriptEngine> painlessScriptEngine = new SetOnce<>();

modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.score.txt → modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.score.txt

@@ -10,6 +10,8 @@
 
 class org.elasticsearch.script.ScoreScript @no_import {
 }
+class org.elasticsearch.script.ScoreScript$Factory @no_import {
+}
 
 static_import {
     double saturation(double, double) from_class org.elasticsearch.script.ScoreScriptUtils

modules/lang-painless/src/test/java/org/elasticsearch/painless/AugmentationTests.java

@@ -26,7 +26,7 @@ public class AugmentationTests extends ScriptTestCase {
     protected Map<ScriptContext<?>, List<Whitelist>> scriptContexts() {
         Map<ScriptContext<?>, List<Whitelist>> contexts = super.scriptContexts();
         List<Whitelist> digestWhitelist = new ArrayList<>(PainlessPlugin.BASE_WHITELISTS);
-        digestWhitelist.add(WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.ingest.txt"));
+        digestWhitelist.add(WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.ingest.txt"));
         contexts.put(DigestTestScript.CONTEXT, digestWhitelist);
 
         return contexts;