[Painless] Add a Map for java names to classes for use in the custom classloader (#34424)

This fixes a bug that wasn't including the class used for a static import 
where only the static import is whitelisted.

Author: jdconrad

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

@@ -96,7 +96,7 @@ public Class<?> findClass(String name) throws ClassNotFoundException {
             if (found != null) {
                 return found;
             }
-            found = painlessLookup.canonicalTypeNameToType(name.replace('$', '.'));
+            found = painlessLookup.javaClassNameToClass(name);
 
             return found != null ? found : super.findClass(name);
         }

modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookup.java

@@ -34,29 +34,39 @@
 
 public final class PainlessLookup {
 
+    private final Map<String, Class<?>> javaClassNamesToClasses;
     private final Map<String, Class<?>> canonicalClassNamesToClasses;
     private final Map<Class<?>, PainlessClass> classesToPainlessClasses;
 
     private final Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods;
     private final Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings;
 
-    PainlessLookup(Map<String, Class<?>> canonicalClassNamesToClasses, Map<Class<?>, PainlessClass> classesToPainlessClasses,
+    PainlessLookup(
+            Map<String, Class<?>> javaClassNamesToClasses,
+            Map<String, Class<?>> canonicalClassNamesToClasses,
+            Map<Class<?>, PainlessClass> classesToPainlessClasses,
             Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods,
             Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings) {
 
+        Objects.requireNonNull(javaClassNamesToClasses);
         Objects.requireNonNull(canonicalClassNamesToClasses);
         Objects.requireNonNull(classesToPainlessClasses);
 
         Objects.requireNonNull(painlessMethodKeysToImportedPainlessMethods);
         Objects.requireNonNull(painlessMethodKeysToPainlessClassBindings);
 
+        this.javaClassNamesToClasses = javaClassNamesToClasses;
         this.canonicalClassNamesToClasses = Collections.unmodifiableMap(canonicalClassNamesToClasses);
         this.classesToPainlessClasses = Collections.unmodifiableMap(classesToPainlessClasses);
 
         this.painlessMethodKeysToImportedPainlessMethods = Collections.unmodifiableMap(painlessMethodKeysToImportedPainlessMethods);
         this.painlessMethodKeysToPainlessClassBindings = Collections.unmodifiableMap(painlessMethodKeysToPainlessClassBindings);
     }
 
+    public Class<?> javaClassNameToClass(String javaClassName) {
+        return javaClassNamesToClasses.get(javaClassName);
+    }
+
     public boolean isValidCanonicalClassName(String canonicalClassName) {
         Objects.requireNonNull(canonicalClassName);
 

modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessLookupBuilder.java

@@ -120,13 +120,23 @@ public static PainlessLookup buildFromWhitelists(List<Whitelist> whitelists) {
         return painlessLookupBuilder.build();
     }
 
+    // javaClassNamesToClasses is all the classes that need to be available to the custom classloader
+    // including classes used as part of imported methods and class bindings but not necessarily whitelisted
+    // individually. The values of javaClassNamesToClasses are a superset of the values of
+    // canonicalClassNamesToClasses.
+    private final Map<String, Class<?>> javaClassNamesToClasses;
+    // canonicalClassNamesToClasses is all the whitelisted classes available in a Painless script including
+    // classes with imported canonical names but does not include classes from imported methods or class
+    // bindings unless also whitelisted separately. The values of canonicalClassNamesToClasses are a subset
+    // of the values of javaClassNamesToClasses.
     private final Map<String, Class<?>> canonicalClassNamesToClasses;
     private final Map<Class<?>, PainlessClassBuilder> classesToPainlessClassBuilders;
 
     private final Map<String, PainlessMethod> painlessMethodKeysToImportedPainlessMethods;
     private final Map<String, PainlessClassBinding> painlessMethodKeysToPainlessClassBindings;
 
     public PainlessLookupBuilder() {
+        javaClassNamesToClasses = new HashMap<>();
         canonicalClassNamesToClasses = new HashMap<>();
         classesToPainlessClassBuilders = new HashMap<>();
 
@@ -189,7 +199,16 @@ public void addPainlessClass(Class<?> clazz, boolean importClassName) {
             throw new IllegalArgumentException("invalid class name [" + canonicalClassName + "]");
         }
 
-        Class<?> existingClass = canonicalClassNamesToClasses.get(canonicalClassName);
+        Class<?> existingClass = javaClassNamesToClasses.get(clazz.getName());
+
+        if (existingClass == null) {
+            javaClassNamesToClasses.put(clazz.getName(), clazz);
+        } else if (existingClass != clazz) {
+            throw new IllegalArgumentException("class [" + canonicalClassName + "] " +
+                    "cannot represent multiple java classes with the same name from different class loaders");
+        }
+
+        existingClass = canonicalClassNamesToClasses.get(canonicalClassName);
 
         if (existingClass != null && existingClass != clazz) {
             throw new IllegalArgumentException("class [" + canonicalClassName + "] " +
@@ -685,6 +704,14 @@ public void addImportedPainlessMethod(Class<?> targetClass, String methodName, C
         }
 
         String targetCanonicalClassName = typeToCanonicalTypeName(targetClass);
+        Class<?> existingTargetClass = javaClassNamesToClasses.get(targetClass.getName());
+
+        if (existingTargetClass == null) {
+            javaClassNamesToClasses.put(targetClass.getName(), targetClass);
+        } else if (existingTargetClass != targetClass) {
+            throw new IllegalArgumentException("class [" + targetCanonicalClassName + "] " +
+                    "cannot represent multiple java classes with the same name from different class loaders");
+        }
 
         if (METHOD_NAME_PATTERN.matcher(methodName).matches() == false) {
             throw new IllegalArgumentException(
@@ -818,6 +845,14 @@ public void addPainlessClassBinding(Class<?> targetClass, String methodName, Cla
         }
 
         String targetCanonicalClassName = typeToCanonicalTypeName(targetClass);
+        Class<?> existingTargetClass = javaClassNamesToClasses.get(targetClass.getName());
+
+        if (existingTargetClass == null) {
+            javaClassNamesToClasses.put(targetClass.getName(), targetClass);
+        } else if (existingTargetClass != targetClass) {
+            throw new IllegalArgumentException("class [" + targetCanonicalClassName + "] " +
+                    "cannot represent multiple java classes with the same name from different class loaders");
+        }
 
         Constructor<?>[] javaConstructors = targetClass.getConstructors();
         Constructor<?> javaConstructor = null;
@@ -952,7 +987,23 @@ public PainlessLookup build() {
             classesToPainlessClasses.put(painlessClassBuilderEntry.getKey(), painlessClassBuilderEntry.getValue().build());
         }
 
-        return new PainlessLookup(canonicalClassNamesToClasses, classesToPainlessClasses,
+        if (javaClassNamesToClasses.values().containsAll(canonicalClassNamesToClasses.values()) == false) {
+            throw new IllegalArgumentException("the values of java class names to classes " +
+                    "must be a superset of the values of canonical class names to classes");
+        }
+
+        if (javaClassNamesToClasses.values().containsAll(classesToPainlessClasses.keySet()) == false) {
+            throw new IllegalArgumentException("the values of java class names to classes " +
+                    "must be a superset of the keys of classes to painless classes");
+        }
+
+        if (canonicalClassNamesToClasses.values().containsAll(classesToPainlessClasses.keySet()) == false ||
+                classesToPainlessClasses.keySet().containsAll(canonicalClassNamesToClasses.values()) == false) {
+            throw new IllegalArgumentException("the values of canonical class names to classes "  +
+                    "must have the same classes as the keys of classes to painless classes");
+        }
+
+        return new PainlessLookup(javaClassNamesToClasses, canonicalClassNamesToClasses, classesToPainlessClasses,
                 painlessMethodKeysToImportedPainlessMethods, painlessMethodKeysToPainlessClassBindings);
     }
 

plugins/examples/painless-whitelist/src/main/java/org/elasticsearch/example/painlesswhitelist/ExampleStaticMethodClass.java

@@ -0,0 +1,26 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.example.painlesswhitelist;
+
+public class ExampleStaticMethodClass {
+    public static int exampleAddInts(int x, int y) {
+        return x + y;
+    }
+}

plugins/examples/painless-whitelist/src/main/resources/org/elasticsearch/example/painlesswhitelist/example_whitelist.txt

@@ -39,4 +39,8 @@ class java.lang.String {
   # existing classes can be "augmented" to have additional methods, which take the object
   # to operate on as the first argument to a static method
   int org.elasticsearch.example.painlesswhitelist.ExampleWhitelistedClass toInt()
+}
+
+static_import {
+  int exampleAddInts(int, int) from_class org.elasticsearch.example.painlesswhitelist.ExampleStaticMethodClass
 }

plugins/examples/painless-whitelist/src/test/resources/rest-api-spec/test/painless_whitelist/30_static.yml

@@ -0,0 +1,26 @@
+# Example test using whitelisted statically imported method
+
+"custom static imported method":
+- do:
+    index:
+      index: test
+      type: test
+      id: 1
+      body: { "num1": 1 }
+- do:
+    indices.refresh: {}
+
+- do:
+    index: test
+    search:
+      body:
+        query:
+          match_all: {}
+        script_fields:
+          sNum1:
+            script:
+              source: "exampleAddInts(2, (int)doc['num1'][0])"
+              lang: painless
+
+- match: { hits.total: 1 }
+- match: { hits.hits.0.fields.sNum1.0: 3 }