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

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,15 @@ public static PainlessLookup buildFromWhitelists(List<Whitelist> whitelists) {
         return painlessLookupBuilder.build();
     }
 
+    private final Map<String, Class<?>> 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 +191,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(), existingClass);
+        } 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 +696,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(), existingTargetClass);
+        } 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 +837,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(), existingTargetClass);
+        } 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 +979,7 @@ public PainlessLookup build() {
             classesToPainlessClasses.put(painlessClassBuilderEntry.getKey(), painlessClassBuilderEntry.getValue().build());
         }
 
-        return new PainlessLookup(canonicalClassNamesToClasses, classesToPainlessClasses,
+        return new PainlessLookup(javaClassNamesToClasses, canonicalClassNamesToClasses, classesToPainlessClasses,
                 painlessMethodKeysToImportedPainlessMethods, painlessMethodKeysToPainlessClassBindings);
     }