Scripting: Fix painless compiler loader to know about context classes

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

@@ -69,17 +69,14 @@ final class Compiler {
     /**
      * A secure class loader used to define Painless scripts.
      */
-    static final class Loader extends SecureClassLoader {
+    final class Loader extends SecureClassLoader {
         private final AtomicInteger lambdaCounter = new AtomicInteger(0);
-        private final PainlessLookup painlessLookup;
 
         /**
          * @param parent The parent ClassLoader.
          */
-        Loader(ClassLoader parent, PainlessLookup painlessLookup) {
+        Loader(ClassLoader parent) {
             super(parent);
-
-            this.painlessLookup = painlessLookup;
         }
 
         /**
@@ -90,6 +87,15 @@ static final class Loader extends SecureClassLoader {
          */
         @Override
         public Class<?> findClass(String name) throws ClassNotFoundException {
+            if (scriptClass.getName().equals(name)) {
+                return scriptClass;
+            }
+            if (factoryClass != null && factoryClass.getName().equals(name)) {
+                return factoryClass;
+            }
+            if (statefulFactoryClass != null && statefulFactoryClass.getName().equals(name)) {
+                return statefulFactoryClass;
+            }
             Class<?> found = painlessLookup.getClassFromBinaryName(name);
 
             return found != null ? found : super.findClass(name);
@@ -139,13 +145,23 @@ int newLambdaIdentifier() {
      * {@link Compiler}'s specified {@link PainlessLookup}.
      */
     public Loader createLoader(ClassLoader parent) {
-        return new Loader(parent, painlessLookup);
+        return new Loader(parent);
     }
 
     /**
-     * The class/interface the script is guaranteed to derive/implement.
+     * The class/interface the script will implement.
+     */
+    private final Class<?> scriptClass;
+
+    /**
+     * The class/interface to create the {@code scriptClass} instance.
+     */
+    private final Class<?> factoryClass;
+
+    /**
+     * An optional class/interface to create the {@code factoryClass} instance.
      */
-    private final Class<?> base;
+    private final Class<?> statefulFactoryClass;
 
     /**
      * The whitelist the script will use.
@@ -154,11 +170,15 @@ public Loader createLoader(ClassLoader parent) {
 
     /**
      * Standard constructor.
-     * @param base The class/interface the script is guaranteed to derive/implement.
+     * @param scriptClass The class/interface the script will implement.
+     * @param factoryClass An optional class/interface to create the {@code scriptClass} instance.
+     * @param statefulFactoryClass An optional class/interface to create the {@code factoryClass} instance.
      * @param painlessLookup The whitelist the script will use.
      */
-    Compiler(Class<?> base, PainlessLookup painlessLookup) {
-        this.base = base;
+    Compiler(Class<?> scriptClass, Class<?> factoryClass, Class<?> statefulFactoryClass, PainlessLookup painlessLookup) {
+        this.scriptClass = scriptClass;
+        this.factoryClass = factoryClass;
+        this.statefulFactoryClass = statefulFactoryClass;
         this.painlessLookup = painlessLookup;
     }
 
@@ -177,7 +197,7 @@ Constructor<?> compile(Loader loader, MainMethodReserved reserved, String name,
                 " plugin if a script longer than this length is a requirement.");
         }
 
-        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, base);
+        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SSource root = Walker.buildPainlessTree(scriptClassInfo, reserved, name, source, settings, painlessLookup,
                 null);
         root.analyze(painlessLookup);
@@ -209,7 +229,7 @@ byte[] compile(String name, String source, CompilerSettings settings, Printer de
                 " plugin if a script longer than this length is a requirement.");
         }
 
-        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, base);
+        ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SSource root = Walker.buildPainlessTree(scriptClassInfo, new MainMethodReserved(), name, source, settings, painlessLookup,
                 debugStream);
         root.analyze(painlessLookup);

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

@@ -36,7 +36,6 @@
 import java.security.PrivilegedAction;
 
 import static java.lang.invoke.MethodHandles.Lookup;
-import static org.elasticsearch.painless.Compiler.Loader;
 import static org.elasticsearch.painless.WriterConstants.CLASS_VERSION;
 import static org.elasticsearch.painless.WriterConstants.CTOR_METHOD_NAME;
 import static org.elasticsearch.painless.WriterConstants.DELEGATE_BOOTSTRAP_HANDLE;
@@ -207,7 +206,7 @@ public static CallSite lambdaBootstrap(
             MethodType delegateMethodType,
             int isDelegateInterface)
             throws LambdaConversionException {
-        Loader loader = (Loader)lookup.lookupClass().getClassLoader();
+        Compiler.Loader loader = (Compiler.Loader)lookup.lookupClass().getClassLoader();
         String lambdaClassName = Type.getInternalName(lookup.lookupClass()) + "$$Lambda" + loader.newLambdaIdentifier();
         Type lambdaClassType = Type.getObjectType(lambdaClassName);
         Type delegateClassType = Type.getObjectType(delegateClassName.replace('.', '/'));
@@ -457,11 +456,11 @@ private static void endLambdaClass(ClassWriter cw) {
     }
 
     /**
-     * Defines the {@link Class} for the lambda class using the same {@link Loader}
+     * Defines the {@link Class} for the lambda class using the same {@link Compiler.Loader}
      * that originally defined the class for the Painless script.
      */
     private static Class<?> createLambdaClass(
-            Loader loader,
+            Compiler.Loader loader,
             ClassWriter cw,
             Type lambdaClassType) {
 

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

@@ -102,10 +102,10 @@ public PainlessScriptEngine(Settings settings, Map<ScriptContext<?>, List<Whitel
         for (Map.Entry<ScriptContext<?>, List<Whitelist>> entry : contexts.entrySet()) {
             ScriptContext<?> context = entry.getKey();
             if (context.instanceClazz.equals(SearchScript.class) || context.instanceClazz.equals(ExecutableScript.class)) {
-                contextsToCompilers.put(context, new Compiler(GenericElasticsearchScript.class,
+                contextsToCompilers.put(context, new Compiler(GenericElasticsearchScript.class, null, null,
                         PainlessLookupBuilder.buildFromWhitelists(entry.getValue())));
             } else {
-                contextsToCompilers.put(context, new Compiler(context.instanceClazz,
+                contextsToCompilers.put(context, new Compiler(context.instanceClazz, context.factoryClazz, context.statefulFactoryClazz,
                         PainlessLookupBuilder.buildFromWhitelists(entry.getValue())));
             }
         }