Update Painless to Allow Augmentation from Any Class (#25360)

Custom whitelists in Painless will need to allow classes to be augmented beyond the currently hard-coded Augmentation class tied to Painless directly. This change allows any class to specify an augmentation on a Painless struct using an appropriate static method. Changes to loading the whitelist have also been created to allow for this specification of a different class for augmentation.

10 files changed, 115 insertions, 60 deletions

diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/Definition.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/Definition.java
index f8bee4e5cf..f0897e7093 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/Definition.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/Definition.java
@@ -199,14 +199,14 @@ public final class Definition {
     public static class Method {
         public final String name;
         public final Struct owner;
-        public final boolean augmentation;
+        public final Class<?> augmentation;
         public final Type rtn;
         public final List<Type> arguments;
         public final org.objectweb.asm.commons.Method method;
         public final int modifiers;
         public final MethodHandle handle;
 
-        public Method(String name, Struct owner, boolean augmentation, Type rtn, List<Type> arguments,
+        public Method(String name, Struct owner, Class<?> augmentation, Type rtn, List<Type> arguments,
                       org.objectweb.asm.commons.Method method, int modifiers, MethodHandle handle) {
             this.name = name;
             this.augmentation = augmentation;
@@ -232,10 +232,10 @@ public final class Definition {
             // otherwise compute it
             final Class<?> params[];
             final Class<?> returnValue;
-            if (augmentation) {
+            if (augmentation != null) {
                 // static method disguised as virtual/interface method
                 params = new Class<?>[1 + arguments.size()];
-                params[0] = Augmentation.class;
+                params[0] = augmentation;
                 for (int i = 0; i < arguments.size(); i++) {
                     params[i + 1] = arguments.get(i).clazz;
                 }
@@ -268,9 +268,9 @@ public final class Definition {
 
         public void write(MethodWriter writer) {
             final org.objectweb.asm.Type type;
-            if (augmentation) {
+            if (augmentation != null) {
                 assert java.lang.reflect.Modifier.isStatic(modifiers);
-                type = WriterConstants.AUGMENTATION_TYPE;
+                type = org.objectweb.asm.Type.getType(augmentation);
             } else {
                 type = owner.type;
             }
@@ -731,7 +731,7 @@ public final class Definition {
                 " with arguments " + Arrays.toString(classes) + ".");
         }
 
-        final Method constructor = new Method(name, owner, false, returnType, Arrays.asList(args), asm, reflect.getModifiers(), handle);
+        final Method constructor = new Method(name, owner, null, returnType, Arrays.asList(args), asm, reflect.getModifiers(), handle);
 
         owner.constructors.put(methodKey, constructor);
     }
@@ -775,10 +775,14 @@ public final class Definition {
                 }
                 addConstructorInternal(className, "<init>", args);
             } else {
-                if (methodName.indexOf("*") >= 0) {
-                    addMethodInternal(className, methodName.substring(0, methodName.length() - 1), true, rtn, args);
+                int index = methodName.lastIndexOf(".");
+
+                if (index >= 0) {
+                    String augmentation = methodName.substring(0, index);
+                    methodName = methodName.substring(index + 1);
+                    addMethodInternal(className, methodName, augmentation, rtn, args);
                 } else {
-                    addMethodInternal(className, methodName, false, rtn, args);
+                    addMethodInternal(className, methodName, null, rtn, args);
                 }
             }
         } else {
@@ -787,8 +791,7 @@ public final class Definition {
         }
     }
 
-    private void addMethodInternal(String struct, String name, boolean augmentation,
-                                   Type rtn, Type[] args) {
+    private void addMethodInternal(String struct, String name, String augmentation, Type rtn, Type[] args) {
         final Struct owner = structsMap.get(struct);
 
         if (owner == null) {
@@ -817,14 +820,20 @@ public final class Definition {
         final Class<?> implClass;
         final Class<?>[] params;
 
-        if (augmentation == false) {
+        if (augmentation == null) {
             implClass = owner.clazz;
             params = new Class<?>[args.length];
             for (int count = 0; count < args.length; ++count) {
                 params[count] = args[count].clazz;
             }
         } else {
-            implClass = Augmentation.class;
+            try {
+                implClass = Class.forName(augmentation);
+            } catch (ClassNotFoundException cnfe) {
+                throw new IllegalArgumentException("Augmentation class [" + augmentation + "]" +
+                    " not found for struct [" + struct + "] using method name [" + name + "].", cnfe);
+            }
+
             params = new Class<?>[args.length + 1];
             params[0] = owner.clazz;
             for (int count = 0; count < args.length; ++count) {
@@ -862,9 +871,10 @@ public final class Definition {
         }
 
         final int modifiers = reflect.getModifiers();
-        final Method method = new Method(name, owner, augmentation, rtn, Arrays.asList(args), asm, modifiers, handle);
+        final Method method =
+            new Method(name, owner, augmentation == null ? null : implClass, rtn, Arrays.asList(args), asm, modifiers, handle);
 
-        if (augmentation == false && java.lang.reflect.Modifier.isStatic(modifiers)) {
+        if (augmentation == null && java.lang.reflect.Modifier.isStatic(modifiers)) {
             owner.staticMethods.put(methodKey, method);
         } else {
             owner.methods.put(methodKey, method);
@@ -966,8 +976,8 @@ public final class Definition {
                             // TODO: we *have* to remove all these public members and use getter methods to encapsulate!
                             final Class<?> impl;
                             final Class<?> arguments[];
-                            if (method.augmentation) {
-                                impl = Augmentation.class;
+                            if (method.augmentation != null) {
+                                impl = method.augmentation;
                                 arguments = new Class<?>[method.arguments.size() + 1];
                                 arguments[0] = method.owner.clazz;
                                 for (int i = 0; i < method.arguments.size(); i++) {
diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/FeatureTestAugmentation.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/FeatureTestAugmentation.java
new file mode 100644
index 0000000000..c1ea19defb
--- /dev/null
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/FeatureTestAugmentation.java
@@ -0,0 +1,32 @@
+/*
+ * 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.painless;
+
+public class FeatureTestAugmentation {
+    public static int getTotal(FeatureTest ft) {
+        return ft.getX() + ft.getY();
+    }
+
+    public static int addToTotal(FeatureTest ft, int add) {
+        return getTotal(ft) + add;
+    }
+
+    private FeatureTestAugmentation() {}
+}
diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/FunctionRef.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/FunctionRef.java
index 6bfe911d97..eb2bb1f554 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/FunctionRef.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/FunctionRef.java
@@ -97,8 +97,8 @@ public class FunctionRef {
         // the Painless$Script class can be inferred if owner is null
         if (delegateMethod.owner == null) {
             delegateClassName = CLASS_NAME;
-        } else if (delegateMethod.augmentation) {
-            delegateClassName = Augmentation.class.getName();
+        } else if (delegateMethod.augmentation != null) {
+            delegateClassName = delegateMethod.augmentation.getName();
         } else {
             delegateClassName = delegateMethod.owner.clazz.getName();
         }
diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFunction.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFunction.java
index 257f2975c9..59b7a333cf 100644
--- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFunction.java
+++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/node/SFunction.java
@@ -135,7 +135,7 @@ public final class SFunction extends AStatement {
 
         org.objectweb.asm.commons.Method method =
             new org.objectweb.asm.commons.Method(name, MethodType.methodType(rtnType.clazz, paramClasses).toMethodDescriptorString());
-        this.method = new Method(name, null, false, rtnType, paramTypes, method, Modifier.STATIC | Modifier.PRIVATE, null);
+        this.method = new Method(name, null, null, rtnType, paramTypes, method, Modifier.STATIC | Modifier.PRIVATE, null);
     }
 
     @Override
diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt
index a1cde1711b..0f86679982 100644
--- a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt
+++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.lang.txt
@@ -36,8 +36,8 @@ class CharSequence -> java.lang.CharSequence {
   IntStream chars()
   IntStream codePoints()
   int length()
-  String replaceAll*(Pattern,Function)
-  String replaceFirst*(Pattern,Function)
+  String org.elasticsearch.painless.api.Augmentation.replaceAll(Pattern,Function)
+  String org.elasticsearch.painless.api.Augmentation.replaceFirst(Pattern,Function)
   CharSequence subSequence(int,int)
   String toString()
 }
@@ -53,17 +53,17 @@ class Iterable -> java.lang.Iterable {
   Iterator iterator()
   Spliterator spliterator()
   # some adaptations of groovy methods
-  boolean any*(Predicate)
-  Collection asCollection*()
-  List asList*()
-  def each*(Consumer)
-  def eachWithIndex*(ObjIntConsumer)
-  boolean every*(Predicate)
-  List findResults*(Function)
-  Map groupBy*(Function)
-  String join*(String)
-  double sum*()
-  double sum*(ToDoubleFunction)
+  boolean org.elasticsearch.painless.api.Augmentation.any(Predicate)
+  Collection org.elasticsearch.painless.api.Augmentation.asCollection()
+  List org.elasticsearch.painless.api.Augmentation.asList()
+  def org.elasticsearch.painless.api.Augmentation.each(Consumer)
+  def org.elasticsearch.painless.api.Augmentation.eachWithIndex(ObjIntConsumer)
+  boolean org.elasticsearch.painless.api.Augmentation.every(Predicate)
+  List org.elasticsearch.painless.api.Augmentation.findResults(Function)
+  Map org.elasticsearch.painless.api.Augmentation.groupBy(Function)
+  String org.elasticsearch.painless.api.Augmentation.join(String)
+  double org.elasticsearch.painless.api.Augmentation.sum()
+  double org.elasticsearch.painless.api.Augmentation.sum(ToDoubleFunction)
 }
 
 # Readable: i/o
@@ -756,8 +756,8 @@ class String -> java.lang.String extends CharSequence,Comparable,Object {
   boolean contentEquals(CharSequence)
   String copyValueOf(char[])
   String copyValueOf(char[],int,int)
-  String decodeBase64*()
-  String encodeBase64*()
+  String org.elasticsearch.painless.api.Augmentation.decodeBase64()
+  String org.elasticsearch.painless.api.Augmentation.encodeBase64()
   boolean endsWith(String)
   boolean equalsIgnoreCase(String)
   String format(Locale,String,def[])
diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.regex.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.regex.txt
index aaea78a7a9..4bf1993528 100644
--- a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.regex.txt
+++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.regex.txt
@@ -42,7 +42,7 @@ class Matcher -> java.util.regex.Matcher extends Object {
   boolean find(int)
   String group()
   String group(int)
-  String namedGroup*(String)
+  String org.elasticsearch.painless.api.Augmentation.namedGroup(String)
   int groupCount()
   boolean hasAnchoringBounds()
   boolean hasTransparentBounds()
diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.txt
index 66f8f67d86..ba50a30042 100644
--- a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.txt
+++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/java.util.txt
@@ -41,13 +41,13 @@ class Collection -> java.util.Collection extends Iterable {
   def[] toArray(def[])
 
   # some adaptations of groovy methods
-  List collect*(Function)
-  def collect*(Collection,Function)
-  def find*(Predicate)
-  List findAll*(Predicate)
-  def findResult*(Function)
-  def findResult*(def,Function)
-  List split*(Predicate)
+  List org.elasticsearch.painless.api.Augmentation.collect(Function)
+  def org.elasticsearch.painless.api.Augmentation.collect(Collection,Function)
+  def org.elasticsearch.painless.api.Augmentation.find(Predicate)
+  List org.elasticsearch.painless.api.Augmentation.findAll(Predicate)
+  def org.elasticsearch.painless.api.Augmentation.findResult(Function)
+  def org.elasticsearch.painless.api.Augmentation.findResult(def,Function)
+  List org.elasticsearch.painless.api.Augmentation.split(Predicate)
 }
 
 class Comparator -> java.util.Comparator {
@@ -123,7 +123,7 @@ class List -> java.util.List extends Collection,Iterable {
   def remove(int)
   void replaceAll(UnaryOperator)
   def set(int,def)
-  int getLength*()
+  int org.elasticsearch.painless.api.Augmentation.getLength()
   void sort(Comparator)
   List subList(int,int)
 }
@@ -163,17 +163,17 @@ class Map -> java.util.Map {
   Collection values()
 
   # some adaptations of groovy methods
-  List collect*(BiFunction)
-  def collect*(Collection,BiFunction)
-  int count*(BiPredicate)
-  def each*(BiConsumer)
-  boolean every*(BiPredicate)
-  Map.Entry find*(BiPredicate)
-  Map findAll*(BiPredicate)
-  def findResult*(BiFunction)
-  def findResult*(def,BiFunction)
-  List findResults*(BiFunction)
-  Map groupBy*(BiFunction)
+  List org.elasticsearch.painless.api.Augmentation.collect(BiFunction)
+  def org.elasticsearch.painless.api.Augmentation.collect(Collection,BiFunction)
+  int org.elasticsearch.painless.api.Augmentation.count(BiPredicate)
+  def org.elasticsearch.painless.api.Augmentation.each(BiConsumer)
+  boolean org.elasticsearch.painless.api.Augmentation.every(BiPredicate)
+  Map.Entry org.elasticsearch.painless.api.Augmentation.find(BiPredicate)
+  Map org.elasticsearch.painless.api.Augmentation.findAll(BiPredicate)
+  def org.elasticsearch.painless.api.Augmentation.findResult(BiFunction)
+  def org.elasticsearch.painless.api.Augmentation.findResult(def,BiFunction)
+  List org.elasticsearch.painless.api.Augmentation.findResults(BiFunction)
+  Map org.elasticsearch.painless.api.Augmentation.groupBy(BiFunction)
 }
 
 class Map.Entry -> java.util.Map$Entry {
diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.txt
index ce78f8a631..94ccc70133 100644
--- a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.txt
+++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.txt
@@ -156,6 +156,8 @@ class org.elasticsearch.painless.FeatureTest -> org.elasticsearch.painless.Featu
   boolean overloadedStatic(boolean)
   Object twoFunctionsOfX(Function,Function)
   void listInput(List)
+  int org.elasticsearch.painless.FeatureTestAugmentation.getTotal()
+  int org.elasticsearch.painless.FeatureTestAugmentation.addToTotal(int)
 }
 
 class org.elasticsearch.search.lookup.FieldLookup -> org.elasticsearch.search.lookup.FieldLookup extends Object {
diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/AugmentationTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/AugmentationTests.java
index acf698e2fc..8618194028 100644
--- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/AugmentationTests.java
+++ b/modules/lang-painless/src/test/java/org/elasticsearch/painless/AugmentationTests.java
@@ -188,4 +188,15 @@ public class AugmentationTests extends ScriptTestCase {
                 exec("Map m = new TreeMap(); m.a = -1; m.b = 1; " +
                      "return m.groupBy((key,value) -> value < 0 ? 'negative' : 'positive')"));
     }
+
+    public void testFeatureTest() {
+        assertEquals(5, exec("org.elasticsearch.painless.FeatureTest ft = new org.elasticsearch.painless.FeatureTest();" +
+            " ft.setX(3); ft.setY(2); return ft.getTotal()"));
+        assertEquals(5, exec("def ft = new org.elasticsearch.painless.FeatureTest();" +
+            " ft.setX(3); ft.setY(2); return ft.getTotal()"));
+        assertEquals(8, exec("org.elasticsearch.painless.FeatureTest ft = new org.elasticsearch.painless.FeatureTest();" +
+            " ft.setX(3); ft.setY(2); return ft.addToTotal(3)"));
+        assertEquals(8, exec("def ft = new org.elasticsearch.painless.FeatureTest();" +
+            " ft.setX(3); ft.setY(2); return ft.addToTotal(3)"));
+    }
 }
diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/PainlessDocGenerator.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/PainlessDocGenerator.java
index 910c4940ab..c29260163c 100644
--- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/PainlessDocGenerator.java
+++ b/modules/lang-painless/src/test/java/org/elasticsearch/painless/PainlessDocGenerator.java
@@ -164,7 +164,7 @@ public class PainlessDocGenerator {
         emitAnchor(stream, method);
         stream.print("]]");
 
-        if (false == method.augmentation && Modifier.isStatic(method.modifiers)) {
+        if (null == method.augmentation && Modifier.isStatic(method.modifiers)) {
             stream.print("static ");
         }
 
@@ -268,12 +268,12 @@ public class PainlessDocGenerator {
         stream.print("link:{");
         stream.print(root);
         stream.print("-javadoc}/");
-        stream.print((method.augmentation ? Augmentation.class : method.owner.clazz).getName().replace('.', '/'));
+        stream.print((method.augmentation != null ? method.augmentation : method.owner.clazz).getName().replace('.', '/'));
         stream.print(".html#");
         stream.print(methodName(method));
         stream.print("%2D");
         boolean first = true;
-        if (method.augmentation) {
+        if (method.augmentation != null) {
             first = false;
             stream.print(method.owner.clazz.getName());
         }
@@ -309,7 +309,7 @@ public class PainlessDocGenerator {
      * Pick the javadoc root for a {@link Method}.
      */
     private static String javadocRoot(Method method) {
-        if (method.augmentation) {
+        if (method.augmentation != null) {
             return "painless";
         }
         return javadocRoot(method.owner);